vulnerability-severity-clf

5-class severity triage for security findings — Info / Low / Medium / High / Critical — over real CVE, EPSS, CISA KEV and OSV package data.

Decoder-style LLMs are the wrong shape for this job: the input is one finding, the output is one label, and you run it a few hundred thousand times a day. This is a ModernBERT encoder with a sequence-classification head — laya's shape — so it scores a finding in milliseconds, on CPU if it has to.

from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch

repo = "NilHRH/vulnerability-severity-clf"
tok = AutoTokenizer.from_pretrained(repo)
model = AutoModelForSequenceClassification.from_pretrained(repo).eval()

text = open("example_finding.txt").read()          # the compact render below
enc = tok(text, truncation=True, max_length=320, return_tensors="pt")
with torch.no_grad():
    logits = model(**enc).logits
print(model.config.id2label[logits.argmax(-1).item()])

Results

Held-out test split, scored once, never used for checkpoint selection:

class	n	precision	recall	F1
Info	163	1.000	0.976	0.988
Low	1072	0.997	0.998	0.998
Medium	1245	0.994	0.995	0.995
High	1094	0.994	0.996	0.995
Critical	24	1.000	0.958	0.979

metric	test
accuracy	99.53%
macro-F1	0.9907

Baselines on the same split

baseline	accuracy	macro-F1
best constant class	0.346	0.103
copy the scanner's own tool severity	0.340	0.193
this model	0.995	0.991

That middle row is the important one. Trivy/Snyk already emit a severity, so the obvious cheat is to echo it back. The policy here disagrees with the scanner's rating on roughly two thirds of findings, so copying buys you 0.34. The model is re-deriving the policy from risk signals, not reading back a field.

GGUF quants - every quant, measured

The sibling repo NilHRH/vulnerability-severity-clf-GGUF holds every quantisation of exactly these weights. Each file was dequantised back into transformers and scored on the same test split, so the accuracy column is measured, not inferred:

quant	size	vs f16	test accuracy	notes
F16	301.1 MB	100%	99.53%	unquantised reference
Q8_0	160.8 MB	53%	99.53%	near-lossless, the safe default
Q6_K	129.3 MB	43%	99.53%	smaller, effectively indistinguishable
Q5_K_M	116.2 MB	39%	99.53%
Q5_K_S	111.1 MB	37%	99.53%
Q5_0	109.9 MB	37%	99.53%
Q4_K_M	106.3 MB	35%	99.56%	best size/quality tradeoff
Q4_K_S	99.5 MB	33%	99.56%
Q4_0	96.0 MB	32%	99.50%
Q3_K_L	100.9 MB	34%	99.56%
Q3_K_M	93.7 MB	31%	99.53%
Q3_K_S	83.9 MB	28%	99.47%
Q2_K	79.9 MB	27%	99.50%	smallest, expect visible damage

The other checkpoint in this repo: custom-head/model.pt

The first model trained on this data used a 2-layer head — CLS → Linear(768→256) → GELU → Linear(256→5). It is the more accurate of the two and it is preserved here untouched, but llama.cpp has no second hidden layer on the modern-bert graph, so it cannot be quantised to GGUF. The weights at the root of this repo are the same encoder re-fitted with the canonical classification head, which is why a GGUF build of this project exists at all.

checkpoint	test accuracy	test macro-F1	GGUF
root: ModernBertForSequenceClassification	99.53%	0.9907	yes, NilHRH/vulnerability-severity-clf-GGUF
custom-head/model.pt	99.56%	0.9836	no

Custom-head results, for completeness:

class	n	precision	recall	F1
Info	163	1.000	0.982	0.991
Low	1072	0.997	1.000	0.999
Medium	1245	0.996	0.995	0.996
High	1094	0.995	0.995	0.995
Critical	24	0.920	0.958	0.939

import torch, torch.nn as nn
from transformers import AutoModel, AutoTokenizer

BASE = "jhu-clsp/ettin-encoder-150m"

class TriageHead(nn.Module):
    def __init__(self, encoder_name, hidden=768, mid=256, n=5, dropout=0.1):
        super().__init__()
        self.encoder = AutoModel.from_pretrained(encoder_name)
        self.drop = nn.Dropout(dropout)
        self.fc1, self.act, self.fc2 = nn.Linear(hidden, mid), nn.GELU(), nn.Linear(mid, n)

    def forward(self, ids, mask):
        h = self.encoder(input_ids=ids, attention_mask=mask).last_hidden_state
        return self.fc2(self.act(self.fc1(self.drop(h[:, 0]))))

tok = AutoTokenizer.from_pretrained(BASE)
model = TriageHead(BASE)
ck = torch.load("custom-head/model.pt", map_location="cpu")
model.load_state_dict(ck.get("state_dict", ck))
model.eval()

Architecture

base encoder	jhu-clsp/ettin-encoder-150m (ModernBERT, English 50k vocab)
params	149.6M total, of which 594k is the classification head
head	CLS pool → Linear(768→768) → GELU → LayerNorm → dropout → Linear(768→5)
input	compact text render of the finding, p50 274 / p95 303 tokens, max_len 320
training	1.0 epoch fine-tune, batch 96, lr 0.0003, bf16, gradient checkpointing
loss	class-weighted cross-entropy (Critical ×12, Info ×4.9) — Critical is 0.8% of rows
hardware	one RTX 5060 Ti 16 GB, 20 min

What it reads

The encoder sees a compact render, not raw JSON:

tool Trivy | scan Container Security | tool severity Critical | category Vulnerability
package libssl3 installed 3.5.4-r0 fix 3.5.5-r0 cwe 787
asset environment Production | product tier tier1 | product internet facing True | ...
business business-impact:critical; compliance:true; data-sensitivity:high; ext-facing:true
cvss CVSS:3.1/AV:N/AC:L/PR:N/UI:N/S:U/C:H/I:H/A:H | epss 0.7078 | kev False | exploited False |
aati 8.13 | exploitability 3.9 | poc 5
av N pr N ac L ui N c H i H a H
title openssl: OpenSSL: Remote code execution or Denial of Service via oversized IV in CMS parsing
description Parsing CMS AuthEnvelopedData message with maliciously crafted AEAD parameters...

Labels are a policy, not human judgement

Gold labels come from a deterministic rule set over CVSS · EPSS · CISA KEV · AATI exploit signals · asset exposure · business context. No LLM is in the labelling loop, and the instruction text lives in the same file as the rules so prompt and label cannot drift.

Consequence worth stating plainly: a high score means the model reproduces that policy. It does not mean it agrees with a human analyst, because no human analyst labelled this data. Swap the policy for your own and retrain — the pipeline is data-in, labels-out.

One rule worth knowing: Critical requires confirmed exploitation (KEV, or EPSS above threshold with a public PoC). CVSS Critical alone caps at High. So a CVSS 9.8 finding with EPSS 0.71 and five PoCs but no KEV is labelled High.

Limitations

• Critical is the thin class. 24 test examples; treat Critical recall as ±1 label wide.
• Validation numbers are selection metrics. The checkpoint was chosen on the split the curves are drawn from. The test figures above are the honest ones.
• Asset context is synthetic. Environment / tier / internet-facing / CIA requirements come from a 48-asset inventory that is internally coherent but invented. Swap in a real inventory and retrain.
• Package overlap across splits. Splits are CVE-disjoint, but the same package and asset recur, so accuracy on a brand-new technology stack will be lower than the headline.
• Not a CVSS calculator. It applies an organisational policy on top of CVSS; it does not reproduce CVSS base scores.

Reproduce

Dataset: 71,107 findings / 60,581 CVEs — 63,974 train / 3,535 val / 3,598 test, CVE-hash split, 100% label re-derivation, 0 contradictory prompts. Trained 2026-10-27.

py -3.13 scripts/00_fetch.py && py -3.13 scripts/02_extract.py
py -3.13 scripts/03_label.py && py -3.13 scripts/04_assemble.py
py -3.13 scripts/verify.py
py -3.13 scripts/daily_run.py --push     # train, score test once, push PT + GGUF