DistilBERT — CVE-LMTune CWE Classifier (Flat)

Part of the CVE-LMTune model suite, a collection of language models fine-tuned for multi-taxonomy vulnerability classification across widely used cybersecurity taxonomies, including CWE, CAPEC, and MITRE ATT&CK.

Paper

Franco Terranova, Sana Rekbi, Abdelkader Lahmadi, Isabelle Chrisment. Multi-Taxonomy Vulnerability Classification with Hierarchically Finetuned Language Models. The 23rd Conference on Detection of Intrusions and Malware & Vulnerability Assessment (DIMVA '26).

Overview

This model performs multi-label CWE classification from vulnerability descriptions. Given a CVE-style description, it predicts one or more CWE identifiers associated with the described vulnerability.

Evaluation Results

The model was evaluated on the held-out test set using standard multi-label classification metrics using sigmoid activation and a default threshold of 0.5.

Ranking Metrics

Threshold = 0.5

Quick Start

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

tokenizer = AutoTokenizer.from_pretrained("Sana9/distilbert-vuln2cwe-flat")
model = AutoModelForSequenceClassification.from_pretrained("Sana9/distilbert-vuln2cwe-flat")

text = "Buffer overflow vulnerability in OpenSSL allows remote attackers to execute arbitrary code."

with torch.no_grad():
    probs = torch.sigmoid(
        model(**tokenizer(text, return_tensors="pt", truncation=True)).logits
    )[0]

predictions = {
    model.config.id2label[i]: p.item()
    for i, p in enumerate(probs)
    if p > 0.5
}

print(predictions)

Citation

@inproceedings{terranova2026multitaxonomy,
  author    = {Franco Terranova and Sana Rekbi and Abdelkader Lahmadi and Isabelle Chrisment},
  title     = {Multi-Taxonomy Vulnerability Classification with Hierarchically Finetuned Language Models},
  booktitle = {Proceedings of the International Conference on Detection of Intrusions and Malware, and Vulnerability Assessment (DIMVA)},
  year      = {2026},
  month     = jul,
  address   = {Chania, Crete, Greece},
  note      = {HAL identifier: hal-05500820v2}
}

Related Resources

• 🤗 Full model suite on Hugging Face
• 💻 CVE-LMTune — Training code (GitHub)
• 📦 Zenodo — Data repository

Disclaimers

• This product is a result of the use of the NVD API but is not endorsed or certified by the NVD. The same for the CVE2CAPEC project and the Hugging Face API.
• This project relies on data publicly available from the CWE, CAPEC, and MITRE ATT&CK projects.
• This work has been partially supported by the French National Research Agency under the France 2030 label (Superviz ANR-22-PECY-0008). The views reflected herein do not necessarily reflect the opinion of the French government.