README.md

metadata

language: en
license: apache-2.0
library_name: pytorch
pipeline_tag: graph-ml
tags:
  - threat-intelligence
  - stix
  - graph-neural-network
  - hgt
  - cybersecurity

STIXBert — A Graph Transformer on native STIX2.1 schema

STIXBert is a Heterogeneous Graph Transformer (HGT) pre-trained with self-supervised objectives on real-world STIX 2.1 threat-intelligence graphs. It produces fixed-size embeddings for every node in a STIX bundle—indicators, malware, attack-patterns, threat-actors, campaigns, and more—enabling downstream tasks such as campaign clustering, ATT&CK technique classification, cross-feed deduplication, infrastructure prediction, and feed quality scoring.

Model Details

Property	Value
Architecture	Heterogeneous Graph Transformer (HGT)
Layers	4
Attention heads	4
Hidden dimension	128
Input dimension	128
Output dimension	128
Dropout	0.1
Parameters	26,916,764 (~102.7 MB fp32)
Text encoder	all-MiniLM-L6-v2 (max 256 tokens)
Framework	PyTorch + PyTorch Geometric

Pre-training Objectives

STIXBert is trained with three complementary self-supervised losses:

1. Masked Node Prediction (weight 1.0) — 15% of node features are zeroed; the model reconstructs them via cosine-similarity loss.
2. Link Prediction (weight 1.0) — binary cross-entropy on existing edges vs. sampled negatives.
3. Temporal Ordering (weight 0.3) — predict which of two nodes appeared first based on STIX timestamps.

Training Procedure

Phase 1 — Hyperparameter Search

• Strategy: Random search (20 trials)
• CV: 5-fold stratified by node_type
• Search epochs per trial: 5
• Swept parameters: hidden_dim, num_heads, num_layers, lr, batch_size, mask_ratio, dropout

Phase 2 — Full Training with Best HPs

• Max epochs: 200 (early stopping patience=30)
• Epochs completed: 164
• Optimizer: ADAMW (weight_decay=0.0001)
• Scheduler: cosine (warmup=10 epochs, min_lr=1e-06)
• Gradient clipping: 1.0
• Mixed precision: Yes
• Class imbalance: weighted

Best Hyperparameters (from Phase 1)

Parameter	Value
batch_size	32
dropout	0.1
hidden_dim	256
lr	0.002
mask_ratio	0.1
num_heads	2
num_layers	6

Final Training Losses (Epoch 164)

Loss	Value
Total	0.4259
Masked Node	0.0576
Link Prediction	0.1791
Temporal Ordering	0.6308

Training Data

Statistic	Value
Total nodes	9,524
Total edges	26,079
Node types	11
Edge types	18

Data sources:

• MITRE ATT&CK — Enterprise, Mobile, ICS (STIX 2.1 bundles from raw.githubusercontent.com/mitre-attack/attack-stix-data)
• ThreatFox — Recent IOCs exported as STIX 2.1 (threatfox.abuse.ch/export/json/recent/)
• DigitalSide Threat-Intel — Community STIX 2.1 bundles (github.com/davidonzo/Threat-Intel)

Node types: attack_pattern, campaign, course_of_action, file, identity, indicator, infrastructure, intrusion_set, malware, tool, vulnerability

Edge types: attack_pattern-[revoked_by]->attack_pattern, attack_pattern-[subtechnique_of]->attack_pattern, campaign-[attributed_to]->intrusion_set, campaign-[uses]->attack_pattern, campaign-[uses]->malware, campaign-[uses]->tool, course_of_action-[mitigates]->attack_pattern, indicator-[indicates]->infrastructure, indicator-[indicates]->malware, infrastructure-[communicates_with]->malware, intrusion_set-[revoked_by]->intrusion_set, intrusion_set-[uses]->attack_pattern, intrusion_set-[uses]->malware, intrusion_set-[uses]->tool, malware-[revoked_by]->malware, malware-[revoked_by]->tool, malware-[uses]->attack_pattern, tool-[uses]->attack_pattern

Intended Uses

Use Case	Description
Campaign clustering	Group malware/indicators by embedding similarity; attribute new IOCs to known campaigns
ATT&CK classification	Fine-tune a linear head to predict MITRE ATT&CK tactics from node embeddings
Cross-feed deduplication	Identify near-duplicate indicators across feeds via cosine similarity
Infrastructure prediction	Predict which infrastructure a malware family will use next
Feed quality scoring	Score feed reliability by measuring embedding alignment with ATT&CK ground truth

Limitations

• Pre-trained on publicly available threat intel only; may not generalize to classified or proprietary feeds without fine-tuning.
• Graph structure depends on relationship quality in source data; missing or incorrect STIX relationships degrade embedding quality.
• Text features are encoded with all-MiniLM-L6-v2 — very long descriptions are truncated to 256 tokens.
• Node types not seen during pre-training will need the graph rebuilt with include_scos=True or additional SDO types.

How to Use

import torch
from huggingface_hub import hf_hub_download

# Download model + config
ckpt_path = hf_hub_download(repo_id='shidey/stixbert', filename='stixbert_best.pt')
cfg_path  = hf_hub_download(repo_id='shidey/stixbert', filename='config.json')
meta_path = hf_hub_download(repo_id='shidey/stixbert', filename='graph_metadata.json')

import json
with open(cfg_path) as f:
    cfg = json.load(f)
with open(meta_path) as f:
    meta = json.load(f)

# Rebuild model (paste STIXBert class or import from your code)
model = STIXBert(
    node_types=meta['node_types'],
    edge_types=[tuple(et) for et in meta['edge_types']],
    input_dim=cfg['model']['input_dim'],
    hidden_dim=meta['best_hyperparameters'].get('hidden_dim', cfg['model']['hidden_dim']),
    num_heads=meta['best_hyperparameters'].get('num_heads', cfg['model']['num_heads']),
    num_layers=meta['best_hyperparameters'].get('num_layers', cfg['model']['num_layers']),
    dropout=meta['best_hyperparameters'].get('dropout', cfg['model']['dropout']),
)
model.load_state_dict(torch.load(ckpt_path, map_location='cpu', weights_only=False))
model.eval()

# Get embeddings
embeddings = model.get_embeddings(x_dict, edge_index_dict)

Citation

If you use STIXBert in your work, please cite:

@software{stixbert2026,
  author = {Dey, Shiladitya},
  title  = {STIXBert: Self-Supervised STIX Graph Foundation Model},
  year   = {2026},
  url    = {https://huggingface.co/shidey/stixbert},
}

Repository

Source code and Colab notebook: github.com/sd1977/STIXBert