Sana9/securebert-vuln2atlas-flat

SecureBERT — CVE-LMTune ATLAS Classifier (Flat)

⚠️ Intended scope — AI-related CVEs only. This model predicts MITRE ATLAS techniques, a taxonomy focused on adversarial threats to AI and machine learning systems. It is intended for analyzing vulnerabilities affecting AI/ML systems, including CVEs reported in resources such as the AVID database. Applying this model to generic, non-AI CVEs falls outside its intended scope and may yield unreliable predictions.

Part of the CVE-LMTune model suite, a collection of fine-tuned language models designed for multi-label vulnerability classification across established cybersecurity taxonomies, including CWE, CAPEC, MITRE ATT&CK, and MITRE ATLAS.

Overview

This model performs multi-label classification against the MITRE ATLAS technique taxonomy using vulnerability descriptions as input. Given a CVE-style description, it predicts one or more associated MITRE ATLAS technique identifiers (AML.T####).

Because publicly available AI-related CVEs annotated with MITRE ATLAS techniques are limited, this model was trained using a transfer-based approach. Specifically, it was fine-tuned on non-AI CVEs associated with MITRE ATT&CK techniques that overlap with MITRE ATLAS techniques. The corresponding ATLAS labels were derived using the ATT&CK–ATLAS mappings provided by the CVE2CAPEC project.

The model was subsequently evaluated on a separate set of real AI-related CVEs annotated with MITRE ATLAS techniques to assess its performance in the intended application domain.

Property	Value
Taxonomy	MITRE ATLAS (Adversarial Threat Landscape for AI Systems) — Techniques
Task	Multi-label text classification
Base model	SecureBERT (ehsanaghaei/SecureBERT)
Input	Vulnerability description (e.g., CVE summary)
Output	One or more MITRE ATLAS technique identifiers
Number of labels	10
Training samples	79,405 non-AI CVEs with overlapping ATT&CK–ATLAS techniques
Evaluation samples	228 real AI-related CVEs annotated with ATLAS techniques
Training and evaluation splits	Non-AI CVEs: 60% training, 20% validation, and 20% testing; AI-related CVEs: separate evaluation set

Labels

The model predicts the following 10 MITRE ATLAS technique labels:

ATLAS ID	Technique
AML.T0006	Active Scanning
AML.T0012	Valid Accounts
AML.T0037	Data from Local System
AML.T0052	Phishing
AML.T0074	Masquerading
AML.T0089	Process Discovery
AML.T0091.000	Application Access Token
AML.T0105	Escape to Host
AML.T0106	Exploitation for Credential Access
AML.T0107	Exploitation for Defense Evasion

Evaluation Results

The model was evaluated on two datasets:

• Non-AI CVEs: A held-out test set drawn from the non-AI CVEs used during training and validation.
• AI-related CVEs: A separate evaluation set of 228 real AI-related CVEs annotated with MITRE ATLAS techniques.

Predictions were generated using sigmoid activation with a default classification threshold of 0.5. Results are reported as mean ± standard deviation across five folds for the non-AI CVEs and five independent runs for the AI-related CVEs.

Classification Metrics

The following table presents the reported classification results. The metric order is Micro-F1, Weighted-F1, Subset Accuracy, and Mean Reciprocal Rank (MRR).

Evaluation dataset	Micro-F1	Weighted-F1	Subset Accuracy	MRR
Non-AI CVEs (test set)	0.823 ± 0.006	0.821 ± 0.008	0.742 ± 0.009	0.885 ± 0.006
AI-related CVEs	0.755 ± 0.013	0.750 ± 0.015	0.708 ± 0.005	0.869 ± 0.009

Interpretation: The results indicate that the model performs well on the held-out non-AI test set and retains measurable predictive performance on the separate set of AI-related CVEs. Performance on the AI-related evaluation set is particularly relevant to the model's intended application. Nevertheless, the limited evaluation sample size and the differences between the training and evaluation domains should be considered when interpreting the results.

Quick Start

The following example demonstrates how to load the model and tokenizer, perform inference on a vulnerability description, and retrieve the predicted ATLAS techniques with probabilities above the default threshold of 0.5.

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

model_name = "Sana9/securebert-vuln2atlas-flat"

tokenizer = AutoTokenizer.from_pretrained(
    model_name,
    use_fast=False,
)

model = AutoModelForSequenceClassification.from_pretrained(model_name)
model.eval()

text = (
    "vLLM Allows Remote Code Execution via PyNcclPipe Communication Service"
)

inputs = tokenizer(
    text,
    return_tensors="pt",
    truncation=True,
    max_length=512,
)

with torch.no_grad():
    logits = model(**inputs).logits
    probabilities = torch.sigmoid(logits)[0]

threshold = 0.5

predictions = {
    model.config.id2label[i]: round(probability.item(), 4)
    for i, probability in enumerate(probabilities)
    if probability.item() >= threshold
}

print("Predicted ATLAS techniques:")
for technique, probability in predictions.items():
    print(f"{technique}: {probability:.4f}")