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

Probleemstelling

Het toewijzen van Common Vulnerabilities and Exposures (CVE) beschrijvingen aan de juiste Common Weakness Enumeration (CWE) categorieën is cruciaal voor het prioriteren en repareren van kwetsbaarheden. De huidige standaard, de National Vulnerability Database (NVD), levert echter vaak onnauwkeurige labels:

• Ruis en inconsistentie: Labels worden door verschillende instanties (CNAs) toegekend met wisselende kwaliteit.
• Te generiek: Vaak worden brede "catch-all" labels gebruikt (bijv. CWE-20 "Improper Input Validation") in plaats van specifieke child-CWE's.
• Schaalbaarheid: Bestaande Large Language Models (LLMs) zoals GPT-4 presteren redelijk, maar zijn duur en computergewicht zwaar. Kleinere modellen presteren vaak aanzienlijk slechter.

Het doel van dit onderzoek is het ontwikkelen van een efficiënt, open-source model dat deze mapping nauwkeurig uitvoert, zonder de rekenkracht van enorme LLMs te vereisen.

Methodologie

1. Dataset Constructie

De auteurs hebben een grote, hoogwaardige dataset samengesteld:

• Bron: 336.777 CVE-entries uit de NVD (1999–2026). Na filtering bleven er 318.979 over.
• Label-verbetering (AI-Refinement): Omdat NVD-labels ruis bevatten, zijn alle beschrijvingen opnieuw gelabeld met Claude Sonnet 4.6.
  • Van de bestaande NVD-labels kwam 73,1% exact overeen met de AI-predicatie.
  • Bij toestaan van parent-child relaties steeg de overeenkomst naar 84,5%.
  • De AI-labels werden gereserveerd als "ground truth" voor het trainingsset, inclusief 74.113 eerder ongelabelde CVE's.
• Splits:
  • Trainingsset (234.770): Alle samples met AI-labels.
  • Validatie/Test (27.780): Enkel samples waarbij de NVD en de AI-labels overeenkwamen. Dit creëert een "high-confidence" testset, maar maakt de evaluatie systematisch makkelijker dan willekeurige NVD-data.
• Decontaminatie: Alle 2.000 CVE's uit het externe CTI-Bench-benchmark zijn verwijderd uit de splits om eerlijke externe evaluatie te garanderen.

2. Model Architectuur en Training

• Basismodel: RoBERTa-base (125M parameters).
• Input: "CVE Description: {tekst}", maximaal 384 tokens.
• Twee-fasen Fine-tuning Strategie:
  1. Fase 1 (Warm-up): De eerste 8 van de 12 transformer-lagen en de embedding-laag zijn bevroren. Alleen de bovenste 4 lagen en de classificatiekop worden getraind (4 epochs, LR $10^{-4}$). Dit stabiliseert de kop zonder de pre-trained representaties te vergeten.
  2. Fase 2 (Full Fine-tuning): Alle parameters worden ontvroren en het volledige model wordt getraind (9 epochs, LR $2 \times 10^{-5}$).
• Hardware: Getraind op een enkele NVIDIA RTX 5080 (16 GB) in ongeveer 4 uur.

Belangrijkste Bijdragen

1. Dataset: Publicatie van een schaalbare CVE-naar-CWE dataset (290K samples) met AI-geoptimaliseerde labels en MITRE ATT&CK annotaties.
2. Efficiëntie: Bewijs dat een klein, taak-specifiek encoder-model (125M parameters) kan concurreren met veel grotere, generieke LLM's (8B+ parameters).
3. Methodologie: Een effectieve twee-fasen fine-tuning aanpak die catastrofale vergeeting voorkomt.
4. Analyse: Een kritische blik op de "granulariteitskloof" in benchmarks, waarbij modellen die specifieke child-CWE's voorspellen onterecht worden bestraft als de ground truth een generieke parent-CWE gebruikt.

Resultaten

Interne Evaluatie (Overeenkomst-gefilterd)

Op de held-out testset (27.780 samples) bereikte het model:

• Top-1 Accuracy: 87,4%
• Macro F1: 60,7% (een verbetering van +15,5% ten opzichte van een TF-IDF baseline).
• Observatie: Het model presteert uitstekend op zeldzame CWE-klassen waar TF-IDF faalt, hoewel de prestaties op zeer zeldzame klassen (<10 voorbeelden) nog steeds laag zijn.

Externe Evaluatie (CTI-Bench)

Op het externe CTI-Bench-benchmark (2.000 CVE's, strikte exacte match):

• RoBERTa (125M): 75,6% nauwkeurigheid (95% CI: 72,8–78,2%).
• Cisco Foundation-Sec-8B-Reasoning (8B): 75,3% nauwkeurigheid.
• GPT-4: 72,0% nauwkeurigheid.
• Conclusie: Het 125M-parameter model is statistisch niet significant verschillend van het 8B-parameter model van Cisco, maar gebruikt 64 keer minder parameters.

Granulariteitsanalyse

Een groot deel van de "fouten" op CTI-Bench is te wijten aan een mismatch in detailniveau:

• Het model voorspelt vaak specifieke child-CWE's (bijv. CWE-121 Stack Buffer Overflow), terwijl de ground truth (NVD) een generieke parent-CWE heeft (bijv. CWE-119 Buffer Overflow).
• Bij een hiërarchie-bewuste evaluatie (waar parent-child overeenkomsten als correct worden geteld), stijgt de score van 75,6% naar 86,5%. Dit suggereert dat het model vaak juistere, specifiekere voorspellingen doet dan de grondwaarheid.

Betekenis en Conclusie

Dit paper demonstreert dat voor gespecialiseerde cybersecurity-taken geen enorme generieke LLM's nodig zijn. Met zorgvuldige datasetconstructie (AI-gefilterde labels) en een slimme trainingstrategie kan een klein, open-source model (RoBERTa-base) prestaties leveren die gelijkwaardig zijn aan state-of-the-art 8B-parameter modellen.

De belangrijkste inzichten zijn:

1. Kostenefficiëntie: Het model is 64x kleiner dan de concurrenten met vergelijkbare prestaties.
2. Kwaliteit van Labels: AI-geoptimaliseerde labels (Claude Sonnet) blijken vaak nauwkeuriger en specifieker dan de originele NVD-labels.
3. Benchmark Beperkingen: Huidige benchmarks straffen modellen af die te specifiek zijn. Een verschuiving naar hiërarchie-bewuste evaluatie is noodzakelijk om de werkelijke waarde van deze modellen te meten.

Alle resources (dataset, model en code) zijn openbaar beschikbaar gesteld via Hugging Face.