Complete the model card: evaluation vs zero-shot baseline, intended use, next steps

README.md

@@ -2,39 +2,131 @@
 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: []
 ---
 
-<!-- This model card has been generated automatically according to the information the Trainer had access to. You
-should probably proofread and complete it, then remove this comment. -->
-
 # vulnerability-attack-technique-classification-roberta-base
 
-This model is a fine-tuned version of [roberta-base](https://huggingface.co/roberta-base) on an unknown dataset.
-It achieves the following results on the evaluation set:
-- Loss: 0.6205
-- F1 Micro: 0.4171
-- F1 Macro: 0.2027
-- Precision Micro: 0.3005
-- Recall Micro: 0.6818
-- Recall At 3: 0.4822
-- Recall At 5: 0.6859
+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.
 
-## Model description
+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*.
 
-More information needed
+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
 
-More information needed
+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
 
-More information needed
+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
 
@@ -49,6 +141,8 @@ The following hyperparameters were used during training:
 - 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 |
@@ -101,3 +195,10 @@ The following hyperparameters were used during training:
 - 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/)