README.md

---
license: gpl-3.0
base_model: FacebookAI/roberta-base
datasets:
  - CIRCL/vulnerability-attack-techniques
language:
  - en
tags:
  - security
  - vulnerability
  - mitre-attack
  - cve
  - bi-encoder
  - sentence-similarity
library_name: transformers
---

# 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`](https://huggingface.co/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`](https://huggingface.co/datasets/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`:

```python
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](https://arxiv.org/abs/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](https://github.com/vulnerability-lookup/VulnTrain)
  as part of the [Vulnerability-Lookup](https://vulnerability.circl.lu) project.