| | --- |
| | tags: [graph-neural-network, cybersecurity, rgcn, pytorch-geometric, multi-cloud] |
| | --- |
| | # 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 |
| | ```python |
| | 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 |
| | ``` |
| | |
