vulnerability-attack-technique-classification-roberta-base

Suggests MITRE ATT&CK (Enterprise) techniques from a free-text vulnerability description. This is a multi-label classifier (roberta-base with a sigmoid head, one output per technique) fine-tuned on CIRCL/vulnerability-attack-techniques, a gold dataset of 1,207 CVEs whose labels come from expert MITRE CTID mappings. Given a CVE description it scores 53 parent techniques (e.g. T1190 Exploit Public-Facing Application, T1505 Server Software Component), for use as a ranked list of candidate techniques for analyst review.

The model is trained with VulnTrain and runs in production on the public Vulnerability-Lookup instance operated by CIRCL, served locally by ML-Gateway: every vulnerability page has an ATT&CK tab with the model's suggestions (example: CVE-2021-44077).

The methodology, evaluation protocol, and the negative result on LLM-assisted label expansion are described in the paper Mapping CVEs to MITRE ATT&CK Techniques: A Curated Gold-Set Classifier and the Limits of LLM-Assisted Label Expansion (arXiv:2607.25572).

DOI: 10.57967/hf/9623

How to use

With VulnTrain, against a live CVE:

vulntrain-infer-attack-classification --cve CVE-2021-44077 \
    --model CIRCL/vulnerability-attack-technique-classification-roberta-base

With plain Transformers:

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

model_id = "CIRCL/vulnerability-attack-technique-classification-roberta-base"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForSequenceClassification.from_pretrained(model_id)
model.eval()

description = (
    "Zoho ManageEngine ServiceDesk Plus before 11306, ServiceDesk Plus MSP "
    "before 10530, and SupportCenter Plus before 11014 are vulnerable to "
    "unauthenticated remote code execution."
)
inputs = tokenizer(description, truncation=True, max_length=512, return_tensors="pt")
with torch.no_grad():
    probs = torch.sigmoid(model(**inputs).logits)[0]

for i in probs.argsort(descending=True)[:5]:
    print(f"{model.config.id2label[int(i)]}  {probs[i]:.4f}")

Technique IDs map to names via the ATT&CK Enterprise matrix; sigmoid >= 0.5 is the prediction threshold used in evaluation, but the model is most useful as a ranker (see the recall@k numbers below).

Intended uses & limitations

Intended: triage assistance — given a vulnerability description, surface a short ranked list of candidate ATT&CK techniques for a human analyst to confirm or reject. This is how it is deployed in Vulnerability-Lookup, where the UI explicitly flags the output as unverified AI-generated suggestions.

Limitations:

• The label space is the 53 parent techniques with at least 5 training examples; sub-techniques are collapsed to their parent, and techniques outside this vocabulary can never be suggested.
• The gold set skews toward exploited-in-the-wild CVEs (CTID's corpus and the KEV catalog), so coverage is best for the techniques common there.
• Trained on 972 CVEs — deliberately label-quality-bound rather than data-bound (see the paper's gold-size scaling curve: every metric still improves monotonically with more curated rows).
• English descriptions only; input is truncated at 512 tokens.
• Scores are not calibrated probabilities.
• Suggestions are not verified mappings; treat them as guidance, never as authoritative CTI.

Training and evaluation data

Labels come from the two public expert sources of CVE→ATT&CK mappings, both following the CTID Mapping ATT&CK to CVE for Impact methodology: the CTID attack_to_cve mappings and the CTID Mappings Explorer KEV mappings (ATT&CK 16.1). Technique IDs revoked since publication are remapped to their successors via the ATT&CK STIX revoked-by relationships. Descriptions are joined from CIRCL/vulnerability-scores (Vulnerability-Lookup). Full details in the dataset card and the VulnTrain documentation.

Splits: 972 train / 106 validation / 118 test examples. Checkpoint selection uses the validation split; the test split is touched once, for the numbers reported here. The paper documents why this matters: selecting the best checkpoint on the test split inflates and destabilizes small-test-split metrics enough to produce spurious conclusions.

Notably, this model was not trained on LLM-generated labels: the paper's controlled experiments show that folding in LLM-labeled CVEs at ≈0.39 agreement with the experts yields no reliable ranking improvement and measurably degrades rare-technique coverage at scale (the comparison checkpoint is published as ...-llm-expanded).

Evaluation

Held-out test split, this checkpoint (seed 42):

Recall@5 = 0.64 means that on average 64% of an unseen CVE's expert-assigned techniques appear in the model's top five suggestions — roughly double the zero-shot embedding-similarity baseline reported in the paper. Across five seeds under the identical protocol, the numbers of record are recall@5 0.673 ± 0.019, recall@3 0.536 ± 0.032, micro-F1 0.410 ± 0.006, macro-F1 0.177 ± 0.014. The complete trainer logs are published in the paper repository.

Training procedure

Binary cross-entropy over 53 sigmoid outputs, with per-label pos_weight balancing (capped at 20) to keep rare techniques trainable. Trained with vulntrain-train-attack-classification (VulnTrain).

Training hyperparameters

The following hyperparameters were used during training:

• learning_rate: 1e-05
• train_batch_size: 32
• eval_batch_size: 32
• seed: 42
• optimizer: Use OptimizerNames.ADAMW_TORCH_FUSED with betas=(0.9,0.999) and epsilon=1e-08 and optimizer_args=No additional optimizer arguments
• lr_scheduler_type: linear
• num_epochs: 40
• max_length: 512
• loss: BCEWithLogitsLoss, balanced pos_weight (min 2.447, max 20.0)
• checkpoint selection: best macro-F1 on the validation split

Training results

Framework versions

• Transformers 5.13.0
• Pytorch 2.12.1+cu130
• Datasets 4.8.5
• Tokenizers 0.22.2

Related artifacts

Citation

@misc{bonhomme2026mappingcvesmitreattck,
      title={Mapping CVEs to MITRE ATT&CK Techniques: A Curated Gold-Set Classifier and the Limits of LLM-Assisted Label Expansion},
      author={Cédric Bonhomme and Alexandre Dulaunoy},
      year={2026},
      eprint={2607.25572},
      archivePrefix={arXiv},
      primaryClass={cs.CR},
      url={https://arxiv.org/abs/2607.25572},
}

Acknowledgements

Developed at CIRCL in the context of the AIPITCH project, co-funded by the European Union.