File size: 2,436 Bytes
0d06661 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | ---
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
```
|