CVE-LMTune: Tuning LMs for Multi-Taxonomy Vulnerability Classification

July 10, 2026 · View on GitHub

📄 Paper Reference

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).

CVE-LMTune is a framework for fine-tuning language models for automated vulnerability classification across multiple cybersecurity taxonomies. The features of the framework include:

🔐 CWE vulnerability weakness classification
⚔️ CAPEC attack pattern classification
🎯 MITRE ATT&CK (sub-)technique classification (Enterprise, ICS, Mobile)
🌳 Hierarchical classification using taxonomy structures to align a DAG of models

Hierarchical Classifier
Hierarchical classification workflow using the MITRE ATT&CK taxonomy.

The framework provides:

Automated vulnerability dataset construction from public cybersecurity repositories
Fine-tuning pipelines for encoder-only language models
Several evaluation strategies for multi-label vulnerability classification
Ready-to-use fine-tuned models available on Hugging Face

🤗 Hugging Face Models

Explore the models already fine-tuned for vulnerability classification released with CVE-LMTune on HuggingFace!

CVE-LMXplore is a web-based interface for exploring the performance of fine-tuned language models on vulnerability classification tasks. It provides interactive visualizations of model predictions, confusion matrices, and feature importance scores. You can also submit custom vulnerability descriptions to compare model predictions across different vulnerability taxonomies.

🌐 Try it online: https://cve-lmxplore.lhs.inria.fr/

Installation

Ensure you have Python and Conda installed, then create the environment:

# Using conda
conda env create -f environment.yml
conda activate cve_lmtune

# Or using pip
pip install -r requirements.txt

Fill in the config/auth.yaml adding your Hugging Face and NVD API keys:

huggingface_key: YOUR_HUGGINGFACE_KEY
nvd_key: YOUR_NVD_KEY

Usage

1. Dataset Construction

Download and preprocess vulnerability data from CVE2CAPEC, NVD, and MITRE:

chmod +x setup_files.sh
./setup_files.sh

This will create the necessary subfolders in the data/ directory. Build the dataset with:

cd data
python3 build_dataset.py

You can optionally sample smaller datasets for experimentation:

python3 sample_data.py -s SIZE \
        [-data DATA_FILE] \
        -c {cwe, capec, tactics, techniques, subtechniques} \
        [--seed 42] [--ml {enterprise, ics, mobile}]

Use --ml enterprise|ics|mobile when classification involves MITRE ATT&CK variants (tactics, techniques, subtechniques). The sampling strategy will make use of the classification level to ensure that samples with no label for that level are not included. If not specified, the default data file will be used (data/data.csv).

Optionally, splits can be computed and stored based on a specific seed for splitting, avoiding recomputation in case of heavy splitting strategies (e.g. stratified KFold on a large dataset):

python3 preprocess_data.py [-d DATA_FILE] \
        -c {cwe, capec, tactics, techniques, subtechniques} \
        [--ml {enterprise, ics, mobile}] \
        --split_method {kfold, train_test_split, temporal, temporal_recursive} \
        [--static_seed] [--seed SEED]

Example:

python3 preprocess_data.py -d data_ics_subtechniques_all.csv \
        -c subtechniques -ml ics \
        --split_method kfold \
        --static_seed --seed SEED

The seed used for splitting can be specified with --seed and should be stored to reload the same split later during fine-tuning (using --data_seed).

2. Fine-Tuning Encoder-Only Models

Run fine-tuning with a chosen encoder-only LM from HuggingFace based on its identifier, on a specific taxonomy, at a specific granularity level, using the following script:

cd ../finetuner
python3 finetune.py [-d DATA_FILE] -m MODEL \
        --cascade_level {0,1,2,...} \
        -c {cwe, capec, tactics, techniques, subtechniques} [--ml {enterprise, ics, mobile}] \
        -sm  {kfold, train_test_split, temporal, temporal_recursive} \
        [--static_seed] [--data_seed DATA_SEED] \
        [--loss_fn {focal_loss, bce, ...}]

When data seed is set, it will be used to search for precomputed split with that seed for that dataset. When --cascade_level is different than 0, hierarchical classification is enabled, and this value specifies until which level of the hierarchy to train. After that level, the hierarchy will be flattened (avoidable using --not_flatten_last_level). This value set to 0 corresponds to the flat classifier.

Example:

python3 finetune.py -d data.csv \
        -m securebert -c cwe -sm train_test_split \
        --static_seed --data_seed 42 \
        --cascade_level 1

Results for each experiment are automatically saved under the finetuner/logs/ directory.

In case of hierarchical classification evaluation, this has to be done joining different trained models in a hierarchy of classification calls, which can be done using:

cd ../model_analysis
python3 test_hierarchy.py -f LOGS_FOLDER_WITH_HIERARCHY --split {train, validation, test}

(Optional) Hyper-parameter Optimization

Optuna is integrated for hyperparameter search of LM and fine-tuning hyper-parameters and can be used as follows.

python3 hyperopt.py ... --num_trials NUM_TRIALS --opt_algo {tpe, random, ...}

The hyperparameters to be explored can be configured in the config/hyperparams_ranges.yaml file. The other arguments are the same as in finetune.py. To monitor the hyperparameter optimization, launch the Optuna dashboard:

optuna-dashboard sqlite:///FILE.db

3. Assessing Text-Generation Models

The repository supports prompt-based classification using text-generation LMs. Currently supported inference providers are Together, OpenAI, and Groq, and additional ones can be easily integrated. To use these models, add the API credentials needed to config/auth.yaml:

...
groq_api_key: YOUR_KEY
openai_api_key: YOUR_KEY
together_api_key: YOUR_KEY
...

The following script queries a selected model across multiple temperature values and prompt variants. All taxonomies are automatically used, and the maximum number of CVEs processed can be controlled via the --threshold argument.

cd ../text_generation
python3 query_model.py -m MODEL \
        --server {groq, openai, together} \
        --temperature_range TEMPERATURE_VALUES \
        --prompt_versions {1,2,3}+ \
        --threshold MAX_CVES

Results of each experiment are saved under the text_generation/results/ directory, organized by model name. Aggregate performance metrics can be computed using:

python3 compute_metrics.py -m MODEL \
        --label_types {enterprise_subtechniques_labels, capec_labels, cwe_labels}+ \
        [--metrics ...]

Advanced Options

The repository offers several advanced features. The cve_stats sub-package allows users to scrape statistics on CVE publications and manual mapping solutions, with each scraping module documented in its own README. Additional scripts in the model_analysis folder enable scalable model evaluation, such as generating score matrices to visualize performance across classification levels. Furthermore, the model_analysis/compute_inference_time script can be used to measure and compare the inference times of different approaches.

This tool can reproduce the results presented in the paper using the same data and configurations. A dedicated Zenodo data repository contains the data leveraged in the paper, as well as the resulting trained models and logs for reproducing the results.

To reproduce the experiments follow the steps outlined in the REPRODUCIBILITY.md file. To reproduce the statistics about CVE disclosures and manual mapping coverage, follow the instructions in the cve_stats/README.md file.

MIT License

The code in this repository is available under the MIT License, see the LICENSE file for more information.

Disclaimers

This repository uses the CVE2CAPEC project, the NVD API, and the Hugging Face API.
This product uses the NVD API but is not endorsed or certified by the NVD.
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.