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

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)

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

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