Complete the model card: evaluation vs zero-shot baseline, intended use, next steps
4bd1d03 verified | 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/) | |