vulnerability-attack-technique-biencoder

A label-semantics bi-encoder that suggests MITRE ATT&CK (Enterprise) techniques for a CVE by scoring the vulnerability description against the official ATT&CK technique descriptions in a shared embedding space. Unlike the companion classification head (CIRCL/vulnerability-attack-technique-classification-roberta-base), it can rank any technique that has an official description — the label is text, not a learned output row.

One shared roberta-base encoder embeds both the CVE text (title + description) and each technique's STIX name+description (citation markup stripped, 256 tokens), mean-pooled and L2-normalized; the score is a learned affine over the cosine. Trained on the curated gold set CIRCL/vulnerability-attack-techniques (~1,200 CVEs, CTID methodology) with per-label-weighted BCE over a 53-parent-technique vocabulary, with VulnTrain (vulntrain-train-attack-biencoder).

When to use which model

• Classification head: best top-5 ranking on the trained vocabulary (recall@5 0.667 ± 0.015 across five seeds).
• This bi-encoder: slightly lower recall@5 (0.643 ± 0.019) but the largest consistent rare-technique gain measured on this task (macro-F1 0.212 ± 0.011 vs 0.176 ± 0.016, +21% relative), and open-vocabulary ranking over all 222 active parent techniques (recall@5 0.515 ± 0.020, 2.3× a generic zero-shot sentence embedder).

Caveat measured in the accompanying paper: zero-shot ranking of techniques absent from training does not benefit from this fine-tuning — in a five-fold label-holdout evaluation the fine-tuned encoder ranked held-out techniques below a generic MiniLM embedder. Rankings for techniques outside the 53-technique training vocabulary should be treated as no better than generic semantic similarity.

Usage

The repository ships technique_texts.json (the exact technique texts used at training time) and the scoring calibration in config.biencoder:

import json, torch
from huggingface_hub import hf_hub_download
from transformers import AutoModel, AutoTokenizer

model_id = "CIRCL/vulnerability-attack-technique-biencoder"
tokenizer = AutoTokenizer.from_pretrained(model_id)
encoder = AutoModel.from_pretrained(model_id).eval()
cfg = encoder.config.biencoder
texts = json.load(open(hf_hub_download(model_id, "technique_texts.json")))

def embed(batch, max_length=512):
    enc = tokenizer(batch, padding=True, truncation=True,
                    max_length=max_length, return_tensors="pt")
    hidden = encoder(**enc).last_hidden_state
    mask = enc["attention_mask"].unsqueeze(-1)
    pooled = (hidden * mask).sum(1) / mask.sum(1)
    return torch.nn.functional.normalize(pooled, dim=-1)

techniques = sorted(texts)
with torch.no_grad():
    technique_emb = embed([texts[t] for t in techniques],
                          cfg["technique_max_length"])
    cve_emb = embed(["Improper neutralization of special elements used "
                     "in an OS command in the web management interface..."])
scores = cfg["logit_scale"] * (cve_emb @ technique_emb.T) + cfg["logit_bias"]
for idx in scores[0].topk(5).indices:
    print(techniques[idx], float(scores[0][idx]))

Evaluation and stratified breakdowns are reproducible with vulntrain-validate-attack-classification --method biencoder --model CIRCL/vulnerability-attack-technique-biencoder (add --candidates full for open-vocabulary ranking over all active parent techniques).

Intended use and limitations

The model generates candidate techniques for analyst review, not authoritative mappings. Technique-to-CVE mapping involves analyst judgment; the training labels inherit the CTID methodology's subjectivity, and the gold set over-represents exploited and enriched CVEs. English descriptions only; parent-level techniques only.

References

• Bonhomme, C., & Dulaunoy, A. (2026). Mapping CVEs to MITRE ATT&CK Techniques: A Curated Gold-Set Classifier and the Limits of LLM-Assisted Label Expansion. arXiv:2607.25572
• Bonhomme, C., & Dulaunoy, A. (2026). Beyond the Description: Structured Metadata and Label Semantics for CVE-to-ATT&CK Mapping. (follow-up paper, in preparation — source of all numbers above)
• Trained with VulnTrain as part of the Vulnerability-Lookup project.