--- library_name: transformers license: mit base_model: roberta-base pipeline_tag: text-classification language: - en datasets: - CIRCL/vulnerability-attack-techniques tags: - vulnerability - cybersecurity - security - cve - mitre-attack - attack-techniques - generated_from_trainer model-index: - name: vulnerability-attack-technique-classification-roberta-base results: [] --- # vulnerability-attack-technique-classification-roberta-base This model suggests **MITRE ATT&CK (Enterprise) techniques** from a vulnerability description. It is a fine-tuned version of [roberta-base](https://huggingface.co/roberta-base) trained on [CIRCL/vulnerability-attack-techniques](https://huggingface.co/datasets/CIRCL/vulnerability-attack-techniques), a dataset of ~1,200 CVEs with analyst-curated technique mappings from the [MITRE Center for Threat-Informed Defense (CTID)](https://ctid.mitre.org/) projects, following the "Mapping ATT&CK to CVE for Impact" methodology. CVSS tells you *how bad* a vulnerability is, CWE *what kind of flaw* it is — ATT&CK tells defenders *what adversary behavior to expect and detect*. This is a **multi-label** classifier (sigmoid head, binary cross-entropy with per-label positive weights): a CVE legitimately maps to several techniques — an exploitation technique (e.g. T1190 *Exploit Public-Facing Application*) plus one or more impacts (e.g. T1059 *Command and Scripting Interpreter*). ## Intended uses & limitations The model **suggests candidate techniques for analyst review**. It must not be treated as an authoritative mapping: a CVE description describes a flaw, while ATT&CK describes attacker behavior around it, and even human annotators disagree on such mappings. Use the top-k suggestions as a triage aid. - **Label space**: 57 parent techniques (sub-techniques collapsed, only techniques with ≥ 5 training examples), enterprise ATT&CK v19.1. Techniques outside this vocabulary can never be suggested. - **Selection bias**: the training data over-represents exploited-in-the-wild vulnerabilities (a large part comes from CISA KEV mappings). - **Small training set** (~1,100 examples): rare-technique performance is weak, as the macro-F1 shows. ## Evaluation Evaluated with `vulntrain-validate-attack-classification` on the dataset's test split, against the zero-shot SMET-style baseline (rank techniques by cosine similarity between the description embedding — `all-MiniLM-L6-v2` — and the official ATT&CK technique descriptions), using the same label vocabulary and metrics: | Metric | Zero-shot similarity baseline | This model | |--------|------------------------------|------------| | recall@1 | 0.118 | **0.220** | | recall@3 | 0.257 | **0.482** | | recall@5 | 0.322 | **0.686** | | recall@10 | 0.491 | **0.842** | | MRR | 0.397 | **0.620** | At the 0.5 threshold: F1 micro 0.417, F1 macro 0.203, precision micro 0.301, recall micro 0.682. The balanced positive weights deliberately favor recall; for the suggestion use case, ranking metrics (recall@k) are the ones that matter — in ~69% of cases the correct techniques appear among the top 5 suggestions, out of 57 candidates. ## Usage ```python 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) description = ( "A Missing Authentication for Critical Function vulnerability in J-Web " "allows an unauthenticated, network-based attacker to upload and " "download arbitrary files and execute commands." ) inputs = tokenizer(description, return_tensors="pt", truncation=True) with torch.no_grad(): probabilities = torch.sigmoid(model(**inputs).logits)[0] top = torch.topk(probabilities, k=5) for probability, index in zip(top.values, top.indices): print(f"{model.config.id2label[index.item()]}: {probability:.2f}") ``` ## Training and evaluation data Trained with [VulnTrain](https://github.com/vulnerability-lookup/VulnTrain): ```bash vulntrain-train-attack-classification \ --base-model roberta-base \ --repo-id CIRCL/vulnerability-attack-technique-classification-roberta-base ``` The dataset provenance (why the labels come from the hand-curated CTID mappings rather than automatically derived CVE→CWE→CAPEC→ATT&CK chains) is documented in the [methodology page](https://github.com/vulnerability-lookup/VulnTrain/blob/main/docs/attack-techniques-dataset.md) and on the [dataset card](https://huggingface.co/datasets/CIRCL/vulnerability-attack-techniques). ## Next steps - **LLM-assisted label expansion**: grow the training set beyond ~1,200 examples by labeling a CWE-stratified CVE sample with the CTID methodology, validated against the analyst gold set before use. This targets the main weakness (rare-technique recall / macro-F1). - Stronger encoders (e.g. security-domain sentence models) as both baseline and base model. - Sub-technique-level labels once the per-label support allows it (`--keep-subtechniques`). - Integration into [Vulnerability-Lookup](https://vulnerability.circl.lu) via [ML-Gateway](https://github.com/vulnerability-lookup/ML-Gateway). ## Training procedure ### Training hyperparameters The following hyperparameters were used during training: - learning_rate: 1e-05 - train_batch_size: 64 - eval_batch_size: 64 - 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 The best checkpoint was selected on validation macro-F1 (epoch 33). ### Training results | Training Loss | Epoch | Step | Validation Loss | F1 Micro | F1 Macro | Precision Micro | Recall Micro | Recall At 3 | Recall At 5 | |:-------------:|:-----:|:----:|:---------------:|:--------:|:--------:|:---------------:|:------------:|:-----------:|:-----------:| | No log | 1.0 | 17 | 0.9141 | 0.1161 | 0.0311 | 0.0672 | 0.4280 | 0.2637 | 0.2959 | | 0.9279 | 2.0 | 34 | 0.8081 | 0.2057 | 0.0305 | 0.1473 | 0.3409 | 0.3201 | 0.3587 | | 0.8472 | 3.0 | 51 | 0.7795 | 0.2369 | 0.0429 | 0.1649 | 0.4205 | 0.3212 | 0.3859 | | 0.7976 | 4.0 | 68 | 0.7677 | 0.2418 | 0.0375 | 0.1752 | 0.3902 | 0.3415 | 0.3950 | | 0.7776 | 5.0 | 85 | 0.7555 | 0.2945 | 0.0616 | 0.2137 | 0.4735 | 0.3449 | 0.4593 | | 0.7742 | 6.0 | 102 | 0.7483 | 0.2943 | 0.0702 | 0.2085 | 0.5 | 0.3530 | 0.4845 | | 0.7742 | 7.0 | 119 | 0.7393 | 0.3159 | 0.0800 | 0.2199 | 0.5606 | 0.3599 | 0.5048 | | 0.7523 | 8.0 | 136 | 0.7281 | 0.3246 | 0.0992 | 0.2212 | 0.6098 | 0.3553 | 0.5116 | | 0.7418 | 9.0 | 153 | 0.7172 | 0.336 | 0.1068 | 0.2283 | 0.6364 | 0.4342 | 0.5668 | | 0.7328 | 10.0 | 170 | 0.7078 | 0.3431 | 0.1230 | 0.2338 | 0.6439 | 0.4063 | 0.5864 | | 0.7166 | 11.0 | 187 | 0.6942 | 0.3546 | 0.1306 | 0.2420 | 0.6629 | 0.4405 | 0.6021 | | 0.6951 | 12.0 | 204 | 0.6858 | 0.3746 | 0.1355 | 0.2635 | 0.6477 | 0.4650 | 0.6256 | | 0.6778 | 13.0 | 221 | 0.6750 | 0.3603 | 0.1339 | 0.2468 | 0.6667 | 0.4252 | 0.6116 | | 0.6778 | 14.0 | 238 | 0.6698 | 0.3661 | 0.1390 | 0.2529 | 0.6629 | 0.4311 | 0.6200 | | 0.6692 | 15.0 | 255 | 0.6623 | 0.3845 | 0.1412 | 0.2684 | 0.6780 | 0.4805 | 0.6796 | | 0.6465 | 16.0 | 272 | 0.6591 | 0.3700 | 0.1510 | 0.2539 | 0.6818 | 0.4287 | 0.6547 | | 0.6499 | 17.0 | 289 | 0.6528 | 0.4036 | 0.1567 | 0.2859 | 0.6856 | 0.4577 | 0.6653 | | 0.6229 | 18.0 | 306 | 0.6505 | 0.4088 | 0.1634 | 0.2886 | 0.7008 | 0.4962 | 0.6785 | | 0.6214 | 19.0 | 323 | 0.6489 | 0.3913 | 0.1485 | 0.2803 | 0.6477 | 0.4661 | 0.6453 | | 0.6092 | 20.0 | 340 | 0.6484 | 0.3825 | 0.1475 | 0.2688 | 0.6629 | 0.4577 | 0.6627 | | 0.6092 | 21.0 | 357 | 0.6409 | 0.4118 | 0.1652 | 0.2935 | 0.6894 | 0.5011 | 0.6852 | | 0.5973 | 22.0 | 374 | 0.6419 | 0.3982 | 0.1585 | 0.2812 | 0.6818 | 0.4647 | 0.6502 | | 0.6046 | 23.0 | 391 | 0.6368 | 0.4095 | 0.1664 | 0.2919 | 0.6856 | 0.4864 | 0.6670 | | 0.5811 | 24.0 | 408 | 0.6395 | 0.3895 | 0.1560 | 0.2785 | 0.6477 | 0.4556 | 0.6572 | | 0.5741 | 25.0 | 425 | 0.6310 | 0.4241 | 0.1740 | 0.3110 | 0.6667 | 0.5179 | 0.6775 | | 0.5697 | 26.0 | 442 | 0.6309 | 0.4014 | 0.1657 | 0.2864 | 0.6705 | 0.4815 | 0.6642 | | 0.5697 | 27.0 | 459 | 0.6292 | 0.4151 | 0.1680 | 0.3014 | 0.6667 | 0.4934 | 0.6880 | | 0.5656 | 28.0 | 476 | 0.6289 | 0.4 | 0.1616 | 0.2901 | 0.6439 | 0.4579 | 0.6607 | | 0.5540 | 29.0 | 493 | 0.6257 | 0.4009 | 0.1798 | 0.2874 | 0.6629 | 0.4710 | 0.6565 | | 0.5506 | 30.0 | 510 | 0.6227 | 0.4110 | 0.1809 | 0.2941 | 0.6818 | 0.4815 | 0.6873 | | 0.5547 | 31.0 | 527 | 0.6242 | 0.4164 | 0.1850 | 0.3012 | 0.6742 | 0.4773 | 0.6747 | | 0.5395 | 32.0 | 544 | 0.6228 | 0.4182 | 0.1844 | 0.3024 | 0.6780 | 0.4661 | 0.6754 | | 0.5407 | 33.0 | 561 | 0.6205 | 0.4171 | 0.2027 | 0.3005 | 0.6818 | 0.4822 | 0.6859 | | 0.5407 | 34.0 | 578 | 0.6187 | 0.4218 | 0.2007 | 0.3038 | 0.6894 | 0.5274 | 0.6971 | | 0.5368 | 35.0 | 595 | 0.6192 | 0.4162 | 0.2018 | 0.2995 | 0.6818 | 0.5025 | 0.6831 | | 0.5342 | 36.0 | 612 | 0.6192 | 0.4176 | 0.1833 | 0.3010 | 0.6818 | 0.5134 | 0.6936 | | 0.5339 | 37.0 | 629 | 0.6205 | 0.4126 | 0.1832 | 0.2980 | 0.6705 | 0.4775 | 0.6859 | | 0.5254 | 38.0 | 646 | 0.6198 | 0.4175 | 0.1851 | 0.3031 | 0.6705 | 0.4817 | 0.6901 | | 0.5318 | 39.0 | 663 | 0.6193 | 0.4159 | 0.1846 | 0.3007 | 0.6742 | 0.4859 | 0.6901 | | 0.5279 | 40.0 | 680 | 0.6191 | 0.4159 | 0.1843 | 0.3007 | 0.6742 | 0.4831 | 0.6859 | ### Framework versions - Transformers 5.13.0 - Pytorch 2.12.1+cu130 - Datasets 4.8.5 - Tokenizers 0.22.2 ## References - [Vulnerability-Lookup](https://vulnerability.circl.lu) — the vulnerability data source - [VulnTrain](https://github.com/vulnerability-lookup/VulnTrain) — training pipeline - [MITRE CTID attack_to_cve](https://github.com/center-for-threat-informed-defense/attack_to_cve) and [Mappings Explorer](https://center-for-threat-informed-defense.github.io/mappings-explorer/) — label sources - MITRE ATT&CK® is a registered trademark of The MITRE Corporation; content used per the [terms of use](https://attack.mitre.org/resources/legal-and-branding/terms-of-use/)