base4security/gnn-threat-scoring

Provenance-GCN — Graph Convolutional Network for APT Detection on Provenance Graphs

Model Description

Provenance-GCN is a lightweight Graph Convolutional Network (GCN) designed to classify OS-level provenance graphs as benign or belonging to specific threat categories. It operates on fixed-size provenance subgraphs (32 nodes) extracted via 2-hop BFS and outputs a classification across 5 threat classes.

The model implements real graph convolution (message passing) following Kipf & Welling (2017), not a simple MLP. The adjacency matrix and node features are passed as a single flattened tensor for ONNX compatibility, and internally reconstructed for graph-aware computation.

Built as the scoring backend for Graph Hunter — a Rust/Tauri threat hunting tool — and compatible with its npu_scorer.rs inference module.

Property	Value
Architecture	2-layer GCN + MLP classifier
Input	Flattened provenance subgraph (batch_size × 1536)
Output	5-class logits (batch_size × 5)
Parameters	6,637 (~6.6K)
File size	46.8 KB
Format	ONNX (opset 18, IR v10)
Producer	PyTorch 2.10.0+cu128
Training hardware	NVIDIA A100-SXM4-80GB (Google Colab)

Architecture

Input [batch, 1536]
  │
  ├── Reshape → Node features X  [batch, 32, 16]
  ├── Reshape → Adjacency matrix A  [batch, 32, 32]
  │
  ├── Â = A + I  (add self-loops)
  ├── A_norm = D⁻¹ Â  (row-normalize)
  │
  ▼
GCN Layer 1:  H¹ = ReLU(A_norm · X · W₁)           (16 → 64)
  │
  ▼
GCN Layer 2:  H² = ReLU(A_norm · H¹ · W₂)          (64 → 32)
  │
  ▼
Global Mean Pooling over 32 nodes                    (32×32 → 32)
  │
  ▼
Classifier:
  Linear(32, 64) → ReLU → Dropout(0.3) → Linear(64, 5)
  │
  ▼
Output [batch, 5]   (logits per class)

Graph Convolution

Each GC layer performs the operation:

$$H^{(l+1)}=\text{ReLU}\text{\hspace{0.17em}⁣}\left({\widehat{D}}^{-1}\widehat{A}{H}^{(l)}{W}^{(l)}\right)H^\{(l+1)\}\; =\; \backslash text\{ReLU\}\backslash !\backslash left(\backslash hat\{D\}^\{-1\}\backslash hat\{A\}\backslash ,\; H^\{(l)\}\backslash ,\; W^\{(l)\}\backslash right)$$

$$\text{where }\widehat{A}=A+I\text{ (adjacency with self-loops) },\widehat{D}\text{ is the degree matrix of }\widehat{A},H^{(l)}\text{ is the node feature matrix at layer l, and }{W}^{(l)}\text{ is the learnable weight matrix.}\backslash text\{where\; \}\backslash hat\{A\}\; =\; A\; +\; I\; \backslash text\{\; (adjacency\; with\; self-loops)\; \},\; \backslash hat\{D\}\; \backslash text\{\; is\; the\; degree\; matrix\; of\; \}\backslash hat\{A\},\; H^\{(l)\}\backslash text\{\; is\; the\; node\; feature\; matrix\; at\; layer\; l,\; and\; \}W^\{(l)\}\backslash text\{\; is\; the\; learnable\; weight\; matrix.\}$$

Weight Tensors

Layer	Shape	Parameters
gc1.weight (via MatMul)	16 × 64	1,024
gc1.bias	64	64
gc2.weight (via MatMul)	64 × 32	2,048
gc2.bias	32	32
classifier.0.weight	32 × 64	2,048
classifier.0.bias	64	64
classifier.3.weight	64 × 5	320
classifier.3.bias	5	5
Total		~6,637

Input Format

The model expects a flattened float32 vector of size 1536 per graph:

Segment	Shape	Description
Node features	32 × 16 = 512	16-dim feature vector per node (see below)
Adjacency matrix	32 × 32 = 1024	Binary adjacency encoding causal edges
Total	1536	Concatenated and flattened

Node Feature Schema (16 dimensions)

Index	Feature	Encoding
0–8	Entity type	One-hot: IP(0), Host(1), User(2), Process(3), File(4), Domain(5), Registry(6), URL(7), Service(8)
9	Out-degree	min(out_degree / 32, 1.0)
10	In-degree	min(in_degree / 32, 1.0)
11–13	Reserved	0.0
14	Total degree	min((in+out) / 64, 1.0)
15	Is center node	1.0 if BFS root, else 0.0

Subgraph Extraction

Subgraphs are extracted as 2-hop BFS neighborhoods from a center node, capped at K_MAX=32 nodes. This matches the extraction logic in gnn_bridge.rs from Graph Hunter.

Output Classes

Index	Class	Description
0	Benign	Normal system activity
1	Exfiltration	Data theft / unauthorized data transfer
2	C2Beacon	Command & Control communication (process → IP patterns)
3	LateralMovement	Lateral movement (process spawning chains)
4	PrivilegeEscalation	Privilege escalation (auth-related anomalies)

Apply softmax to logits for probabilities. Threat score = 1.0 - P(Benign).

Training Details

Dataset: StreamSpot

The model was trained on StreamSpot (Manzoor et al., 2016), a benchmark dataset of OS-level provenance graphs:

Property	Value
Total graphs	600
Benign graphs	500 (IDs 0–299, 400–599)
Attack graphs	100 (IDs 300–399, drive-by-download scenario)
Raw edges	89,770,902
Parsed nodes	5,046,878
Parsed edges	7,638,242
Node types	Process, File, IP
Malicious nodes	889,080

Subgraph Sampling

From the full graph, 1,000,000 subgraphs were extracted (balanced 500K malicious + 500K benign) using 2-hop BFS with K_MAX=32. Threat labels were assigned via heuristics on the attack subgraph structure (network fan-out → C2Beacon, process spawning → LateralMovement, auth edges → PrivilegeEscalation, default malicious → Exfiltration).

Training Configuration

Hyperparameter	Value
Train / Val / Test split	68% / 12% / 20% (stratified)
Batch size	64
Optimizer	Adam (lr=0.001, weight_decay=1e-4)
Scheduler	ReduceLROnPlateau (factor=0.5, patience=5)
Loss	CrossEntropyLoss with inverse-frequency class weights
Max epochs	100
Early stopping	Patience = 10 (on validation loss)
Dropout	0.3 (classifier only)
Hardware	NVIDIA A100-SXM4-80GB
Export	torch.onnx.export (opset 18, dynamic batch axis)

Usage

Python (ONNX Runtime)

import numpy as np
import onnxruntime as ort

session = ort.InferenceSession("provenance_gcn.onnx")

# Build input: 32 nodes × 16 features + 32×32 adjacency, flattened
graph_input = np.zeros((1, 1536), dtype=np.float32)

# Example: center node is a Process (index 3) with high out-degree
graph_input[0, 0 * 16 + 3] = 1.0   # node 0 = Process
graph_input[0, 0 * 16 + 9] = 0.25  # out-degree
graph_input[0, 0 * 16 + 15] = 1.0  # center node flag

# Add neighbor IPs and edges in adjacency block (offset 512)
for i in range(1, 6):
    graph_input[0, i * 16 + 0] = 1.0          # node i = IP
    graph_input[0, 512 + 0 * 32 + i] = 1.0    # edge: node 0 → node i

# Run inference
logits = session.run(None, {"input": graph_input})[0]
probs = np.exp(logits) / np.exp(logits).sum(axis=1, keepdims=True)
threat_score = 1.0 - probs[0, 0]  # 1 - P(Benign)

CLASSES = ["Benign", "Exfiltration", "C2Beacon", "LateralMovement", "PrivilegeEscalation"]
print(f"Prediction: {CLASSES[np.argmax(logits)]}, Threat Score: {threat_score:.3f}")

Batch Inference

batch = np.random.randn(256, 1536).astype(np.float32)
logits = session.run(None, {"input": batch})[0]
predictions = np.argmax(logits, axis=1)

Rust (Graph Hunter / ort crate)

use ort::{Session, Value};
use ndarray::Array2;

let session = Session::builder()?.commit_from_file("provenance_gcn.onnx")?;
let input = Array2::<f32>::zeros((1, 1536));
let outputs = session.run(ort::inputs!["input" => input]?)?;
let logits = outputs[0].extract_tensor::<f32>()?;

The ONNX interface is identical to what npu_scorer.rs expects: 1536 floats in, 5 logits out. Drop-in compatible with Graph Hunter.

Intended Use

• Primary: Real-time classification of OS provenance subgraphs for threat detection in SOC/EDR pipelines, as the scoring backend for Graph Hunter.
• Secondary: Research on graph-based threat detection, GNN benchmarking for cybersecurity, edge/NPU deployment scenarios (model fits comfortably on NPUs like Hailo-8L).
• Out of scope: This model is not a standalone security product. It should be integrated as one signal within a broader detection pipeline.

Limitations

• Fixed graph size: 32 nodes max — larger provenance graphs must be windowed via k-hop BFS before inference.
• Heuristic labels: Threat subcategories (Exfiltration, C2, LateralMovement, PrivEsc) are assigned via structural heuristics on StreamSpot attack graphs, not from ground-truth APT campaign labels.
• Single attack type: StreamSpot only contains one attack scenario (drive-by-download). Generalization to other APT campaigns (APT29, APT3, etc.) requires retraining on datasets like DARPA TC E3 CADETS/Theia.
• Lightweight by design: 6.6K parameters enables edge deployment but limits capacity for complex multi-stage attack patterns.
• No adversarial evaluation: Not tested against evasion techniques targeting GNN-based detectors.

References

• Kipf, T.N. & Welling, M. (2017). Semi-Supervised Classification with Graph Convolutional Networks. ICLR 2017.
• Manzoor, E., Milajerdi, S.M. & Akoglu, L. (2016). Fast Memory-efficient Anomaly Detection in Streaming Heterogeneous Graphs. KDD 2016. (StreamSpot dataset)
• DARPA Transparent Computing Program — TC Engagement 3 datasets

Citation

@misc{provenance-gcn-2025,
  title   = {Provenance-GCN: Lightweight Graph Convolutional Network for APT Detection on Provenance Graphs},
  author  = {BASE4 Security R\&D+i},
  year    = {2026},
  note    = {Trained on StreamSpot, exported as ONNX for Graph Hunter}
}

Contact

• Organization: BASE4 Security — R&D+i Division
• Location: Buenos Aires, Argentina