README.md

---
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

```python
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:

```bibtex
@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](https://github.com/sd1977/STIXBert)