adarsh-aur/rgcn-security-zero-embedding

Stage 6 — Heterogeneous Structural GNN (RGCN)

Multi-Cloud Threat Detection Pipeline — Holistic Version

Library: PyTorch Geometric (RGCNConv) Design: Schema-agnostic — works with any node/edge types, any feature dimension

Architecture

• Input: any fdim → zero-padded to MAX_FDIM=1024 → Linear(1024→256)
• 3-layer RGCNConv (256→256→128), num_relations=20
• PEFT Adapters (rank=16) after layers 1 and 2
• DistMult edge anomaly scoring per relation type

Master Schema

• Node types (15): User, VM, IP, Role, CVE, Container, CloudAccount, Subnet, Bucket, Function, Cluster, Pod, Database, LoadBalancer, Gateway
• Edge types (20): EXPLOITS, CROSS_CLOUD_ACCESS, CONNECTS_TO, EXPLOITS, ACCESS, ASSUMES_ROLE, CONNECTS_TO, RESTART_VM, START_VM, STOP_VM, DEPLOYED_ON, HAS_VULN, LATERAL_MOVEMENT, GRANTS_ACCESS, RUNS_ON, EXPLOITS, ACCESS, TRIGGERS, ACCESS, CONTAINS

Active Schema (this run)

• Node types: User, VM, IP, Role, CVE
• Edge types: 10
• Supervised: User, Role

Ablation Results (Test Set)

Model	Type	Params	User_AUC	User_F1	User_AP	Role_AUC	Role_F1	Role_AP
RGCN	PRIMARY	8,515,599	0.5	0	0.5	0.5	0	0.5
GCN	BASELINE	5,219,855	0.5	0	0.5	0.5	0	0.5
GAT	BASELINE	5,320,207	0.5	0	0.5	0.5	0	0.5
SAGE	BASELINE	5,383,695	0.5	0	0.5	0.5	0	0.5

Usage — Stage 7 API Integration

import torch
from huggingface_hub import hf_hub_download

# Load once, call forever
ckpt  = torch.load(hf_hub_download("adarsh-aur/rgcn-security-zero-embedding", "model_RGCN.pt"))
model = HeteroRGCN()
model.load_state_dict(ckpt['model_state_dict'])
model.eval()

# Works with any graph snapshot from Stage 5
with torch.no_grad():
    h_v, offsets, logits = model(graph_snapshot)
    # h_v shape: [total_nodes, 128]  → feed to Stage 7 GRU
    # New node/edge types: silently skipped
    # Missing node/edge types: silently skipped
    # Different fdim: auto-padded/truncated to 1024