💬 NLP

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

This paper presents a 125M-parameter fine-tuned RoBERTa model that achieves competitive CVE-to-CWE classification performance with significantly larger LLMs, leveraging a large-scale, AI-refined dataset to substantially outperform traditional baselines on rare weakness categories.

Original authors: Nikita Mosievskiy

Published 2026-03-17

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Nikita Mosievskiy

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a librarian in a massive, chaotic library called the Vulnerability Database. Every day, thousands of new books (called CVEs) arrive describing security flaws in software. Your job is to sort these books onto the correct shelves. The shelves are labeled with CWEs (Common Weakness Enumerations), which are categories like "Broken Lock," "Open Window," or "Weak Password."

The problem? The original labels on the books are messy. Sometimes the librarian who wrote the description used a vague label like "Something is broken" (a generic parent category) instead of the specific "The lock is rusted" (a specific child category). This makes it hard to find the right fix quickly.

Enter Nikita Mosievskiy, a researcher who decided to build a super-smart, but tiny, robot librarian to help sort these books. Here is the story of how they did it, using simple analogies.

1. The Goal: A Tiny Robot vs. Giant Brains

Usually, to solve a hard puzzle like this, people use Giant Brains (Huge AI models like GPT-4 or Cisco's 8B model). These are like massive, all-knowing libraries that can read anything but are heavy, expensive, and slow to run.

Mosievskiy asked: "Can we build a tiny, specialized robot that is just as good at sorting these specific books, but is 64 times smaller and cheaper?"

The answer was Yes. They built a robot based on RoBERTa, a model with only 125 million parameters (think of this as the robot's "brain size"). It's tiny compared to the giants, but it learned to be incredibly efficient.

2. The Secret Sauce: Cleaning the Books

The biggest problem wasn't the robot; it was the messy books. The original labels were inconsistent.

The Old Way: Just use the messy labels as they are.
The New Way: The researcher hired a super-smart AI assistant (Claude Sonnet) to read every single book and rewrite the labels to be precise.
- Analogy: Imagine the original label said "The door is broken." The AI assistant read the description and changed it to "The door hinge is snapped."
- They did this for nearly 300,000 books. They even checked to make sure the AI assistant and the original librarians agreed on the easy ones, creating a "Gold Standard" test set.

3. The Training: "Warm-up" then "Full Sprint"

Training a robot to be perfect usually involves two steps, which the researcher called Two-Phase Training:

Phase 1 (The Warm-up): Imagine teaching a student to write an essay. First, you let them practice just the conclusion and the title, while keeping their main knowledge base (the first 8 layers of the brain) frozen. This prevents them from forgetting what they already know.
Phase 2 (The Full Sprint): Once the student knows how to structure the answer, you let them rewrite the whole essay, adjusting every part of their brain to fit the new task.

This method helped the tiny robot learn faster and better than if they had just thrown it into the deep end immediately.

4. The Results: The Underdog Wins

The researchers put their tiny robot to the test against the giants on a standardized exam called CTI-Bench.

The Giants: Cisco's massive 8-billion-parameter model scored 75.3%.
The Tiny Robot: Mosievskiy's 125-million-parameter model scored 75.6%.

The Result: The tiny robot beat the giant, even though the giant is 64 times bigger! They are statistically tied, but the tiny robot is much faster and cheaper to run.

5. The "Gotcha": The Hierarchy Trap

There was a catch. The exam (CTI-Bench) sometimes used the old, vague labels (e.g., "Broken Door"), while the robot was trained to give specific answers (e.g., "Broken Hinge").

If the robot said "Broken Hinge" and the exam key said "Broken Door," the exam marked it wrong.
The Reality: The robot was actually more helpful! Knowing the hinge is broken is better than just knowing the door is broken.
When the researchers adjusted the grading to accept "specific" answers as correct for "general" categories, the robot's score jumped from 75.6% to 86.5%.

6. Why This Matters

This paper is like proving you don't need a Ferrari to win a race on a specific track; a well-tuned go-kart can do the job just as well if the track is set up correctly.

For Security Teams: You don't need to buy expensive, massive AI systems to sort vulnerabilities. A small, open-source model can do the heavy lifting.
For the Industry: It highlights that the "gold standard" labels we've been using for years are often too vague. We need to be more specific to fix security holes effectively.

In short: The researcher built a tiny, specialized tool, taught it with super-clean data, and showed it could outperform (or tie) the massive, expensive giants in the field of cybersecurity, all while being much more efficient. And best of all, they shared the blueprints (the code and data) with everyone for free.

1. Problem Statement

The mapping of Common Vulnerabilities and Exposures (CVE) descriptions to Common Weakness Enumeration (CWE) categories is critical for vulnerability triage and remediation. However, current automated and manual labeling processes face significant challenges:

Noisy Ground Truth: Labels assigned by the National Vulnerability Database (NVD) are often inconsistent, too generic (e.g., using "Improper Input Validation" as a catch-all), or incorrect.
Model Inefficiency: Recent Large Language Models (LLMs) like GPT-4 and specialized 8B-parameter models (e.g., Cisco's Foundation-Sec-8B) have shown promise but require massive computational resources.
Benchmark Limitations: Existing benchmarks often suffer from data leakage (random splits of similar CVEs) or strict evaluation metrics that penalize models for predicting specific child CWEs when the ground truth uses generic parent CWEs.

2. Methodology

A. Dataset Construction

The authors constructed a large-scale dataset of 290,000 samples (234,770 for training, 27,780 for testing) from the NVD (1999–2026).

AI-Refined Labeling: Instead of relying solely on NVD labels, the authors used Claude Sonnet 4.6 to relabel all 318,979 CVE descriptions.
- Agreement Filtering: The test set consists only of samples where the NVD label and the AI label agree, creating a high-confidence "easy" subset.
- Training Data: Includes all samples using the AI-refined labels, including cases where NVD and AI disagreed (prioritizing the AI label) and previously unlabeled CVEs.
Decontamination: All 2,000 samples from the external CTI-Bench benchmark were removed from the training data to ensure a clean external evaluation.
Additional Annotations: The dataset includes MITRE ATT&CK technique mappings for future multi-task research.

B. Model Architecture & Training

Base Model: RoBERTa-base (125M parameters) with a classification head for 205 CWE classes.
Input: CVE descriptions formatted as "CVE Description: {text}", truncated to 384 tokens.
Two-Phase Fine-Tuning Strategy:
1. Phase 1 (Warm-up): Freeze the first 8 of 12 transformer layers and embeddings. Train only the top 4 layers and the classification head for 4 epochs. This stabilizes the classifier head without disrupting pretrained representations.
2. Phase 2 (Full Fine-tuning): Unfreeze all parameters and train for 9 epochs with a lower learning rate ( $2 \times 10^{-5}$ ).
Hyperparameters: Batch size 32, Label smoothing 0.1, bf16 precision, trained on a single NVIDIA RTX 5080 (~4 hours total).

3. Key Contributions

High-Quality Dataset: Release of a 290K-sample CVE-to-CWE dataset with AI-refined labels and MITRE ATT&CK mappings.
Efficient Architecture: Demonstration that a 125M parameter task-specific encoder can match the performance of 8B parameter general-purpose LLMs.
Training Innovation: Validation of a two-phase fine-tuning approach that prevents catastrophic forgetting and improves convergence.
Benchmark Analysis: Identification of CWE hierarchy granularity as a major confounding factor in current benchmarks, arguing that predicting specific child CWEs should not be penalized when ground truth uses generic parent CWEs.

4. Results

Internal Evaluation (Agreement-Filtered Test Set)

Top-1 Accuracy: 87.4% (vs. 84.9% for a TF-IDF baseline).
Macro F1: 60.7% (vs. 45.2% for TF-IDF).
- Significance: The large gap in Macro F1 (+15.5 points) indicates the model significantly outperforms keyword-based baselines on rare CWE classes, whereas the TF-IDF baseline performs well only on common classes (e.g., XSS, SQL Injection).

External Evaluation (CTI-Bench)

Evaluated on the CTI-RCM task (strict exact CWE ID match) with 2,000 zero-shot samples.

RoBERTa-base (125M): 75.6% accuracy (95% CI: 72.8–78.2%).
Cisco Foundation-Sec-8B-Reasoning (8B): 75.3% accuracy.
GPT-4: 72.0%.
Comparison: The 125M model is statistically indistinguishable from the 8B Cisco model while using 64× fewer parameters. It outperforms other open-weight models like LLaMA3-8B (44.7%) and LLaMA3-70B (65.9%).

Hierarchy-Aware Analysis

When evaluating with "hierarchy-aware" metrics (where a specific child CWE prediction counts as correct if it maps to the generic parent CWE in the ground truth):

The model's score on CTI-Bench jumps from 75.6% to 86.5%.
This suggests that many "errors" are actually superior predictions (more specific) that are penalized by the strict benchmark definition.

5. Significance and Conclusion

Parameter Efficiency: The study proves that for specific NLP tasks like vulnerability classification, highly specialized, smaller models can outperform or match massive general-purpose LLMs if trained on high-quality, curated data.
Data Quality over Model Size: The success is attributed primarily to the AI-refined labels (which are shown via manual review to be more accurate/specific than raw NVD labels) and the two-phase training strategy, rather than just model architecture.
Benchmark Critique: The paper highlights that current benchmarks may unfairly penalize models that provide more granular, useful security insights.
Resource Availability: The authors have open-sourced the dataset, the trained model, and the training code, facilitating further research in automated vulnerability analysis.

Limitations: The model is English-only, relies on single-label classification (some CVEs have multiple weaknesses), and the evaluation relies on AI-generated labels which, while validated, introduce a potential circularity bias.