CVE-RiskRank Fusion

Late-fusion exploitation-priority model for CVEs. Predicts the calibrated probability that a vulnerability becomes an operational priority (CISA KEV membership as the exploitation-priority proxy), instead of ranking by CVSS base score alone.

Architecture

Component	Detail
Text encoder	frozen `sentence-transformers/all-MiniLM-L6-v2` (384-d, normalized)
Tabular branch	StandardScaler + one-hot encoding over CVSS vector components, CWE, CPE footprint, reference signals (122-d)
Fusion head	MLP 506 → 256 → 64 → 1, ReLU + dropout, class-weighted BCE
Calibration	isotonic regression fitted on the temporal validation split only

Trained on alirezaaminzadeh/cve-riskrank with a strict temporal split: train ≤ 2024-12-31, validation 2025 H1, test 2025 H2. Early stopping on validation PR-AUC. All training ran on Hugging Face ZeroGPU hardware.

Results (test split, 2025 H2, 91 positives / 31,408 CVEs)

Metric	Fusion (calibrated)	CVSS-score ranking
PR-AUC	0.119	0.012
ROC-AUC	0.946	—
Recall@100	0.220	—
Recall@500	0.473	—
NDCG@100	0.243	—
Brier score	0.0028	—
Expected calibration error	0.0019	—

Raw (uncalibrated) test PR-AUC: 0.135. See metrics.json for the full validation/test breakdown and the baseline comparison in cve-riskrank-baseline.

Files

File	Content
`fusion_head.pt`	fusion MLP state dict
`tabular_preprocessor.joblib`	fitted scikit-learn ColumnTransformer
`calibrator.joblib`	fitted isotonic regression
`config.json`	dimensions and feature column lists
`metrics.json`	full evaluation results
`test_predictions.parquet`	calibrated scores for the temporal test split

Usage

import json

import joblib
import numpy as np
import pandas as pd
import torch
from huggingface_hub import hf_hub_download
from sentence_transformers import SentenceTransformer

REPO = "alirezaaminzadeh/cve-riskrank-fusion"
cfg = json.load(open(hf_hub_download(REPO, "config.json")))
pre = joblib.load(hf_hub_download(REPO, "tabular_preprocessor.joblib"))
cal = joblib.load(hf_hub_download(REPO, "calibrator.joblib"))

class FusionHead(torch.nn.Module):
    def __init__(self, dim_in):
        super().__init__()
        self.net = torch.nn.Sequential(
            torch.nn.Linear(dim_in, 256), torch.nn.ReLU(), torch.nn.Dropout(0.2),
            torch.nn.Linear(256, 64), torch.nn.ReLU(), torch.nn.Dropout(0.1),
            torch.nn.Linear(64, 1))
    def forward(self, x):
        return self.net(x).squeeze(-1)

head = FusionHead(cfg["text_dim"] + cfg["tabular_dim"])
head.load_state_dict(torch.load(hf_hub_download(REPO, "fusion_head.pt"),
                                map_location="cpu", weights_only=True))
head.eval()
encoder = SentenceTransformer(cfg["encoder"])

# `row` follows the cve-riskrank dataset schema
row = {...}
emb = encoder.encode([row["description"]], normalize_embeddings=True)
tab = pre.transform(pd.DataFrame([row])).astype(np.float32)
x = torch.from_numpy(np.hstack([emb, tab]).astype(np.float32))
with torch.no_grad():
    prob = float(cal.predict([torch.sigmoid(head(x)).item()])[0])
print(f"exploitation priority: {prob:.1%}")

Intended use and limitations

Designed for patch prioritization support in vulnerability management workflows. KEV membership is a proxy label: it is curated, US-centric, and lags disclosure, so recent CVEs carry conservative labels. Scores should complement — not replace — asset-context-aware risk assessment.

Demo

CVE-RiskRank Space — live NVD triage, agent advisories, ranking explorer, and evaluation dashboard.

License

MIT

alirezaaminzadeh/cve-riskrank-fusion