| | --- |
| | language: en |
| | license: mit |
| | tags: |
| | - graph-neural-network |
| | - rgcn |
| | - heterogeneous-graph |
| | - cybersecurity |
| | - anomaly-detection |
| | - pytorch-geometric |
| | - multi-cloud |
| | - threat-detection |
| | pipeline_tag: graph-ml |
| | --- |
| | |
| | # Stage 6: Heterogeneous RGCN — Multi-Cloud Threat Detection |
| |
|
| | Part of the research project: |
| | **Intelligent Threat Detection using GNN in Multi-Cloud Environments** |
| |
|
| | ## Architecture |
| |
|
| | | Property | Value | |
| | |---|---| |
| | | Model | Heterogeneous RGCN (RGCNConv) | |
| | | Framework | PyTorch Geometric | |
| | | Init Strategy | GraphSAGE-style projection (Stanford) | |
| | | PEFT | GNN Adapter layers (residual bottleneck) | |
| | | Hidden Dim | 256 | |
| | | Output Dim | 128 | |
| | | Layers | 3 | |
| | | Node Types | user, device, cve, cloud, role, ip | |
| | | Edge Types | assumes_role, accesses, connects_to, has_vuln, deployed_on, lateral_movement, exploits, grants_access | |
| | | Task | Node threat scoring + edge anomaly detection | |
| | | Best Val AUC | 0.9853 | |
| |
|
| | ## Graph Schema |
| |
|
| | ``` |
| | user ──assumes_role──► role ──grants_access──► cloud |
| | user ──accesses──────► device ──has_vuln──────► cve |
| | user ──connects_to───► ip ──connects_to───────► device |
| | device ──lateral_movement──► device |
| | device ──deployed_on──► cloud |
| | cve ──exploits──────────────► device |
| | ``` |
| |
|
| | ## Sample Output |
| |
|
| | ``` |
| | Node: user_4456 |
| | Type: user |
| | Probability: 0.7231 |
| | Risk: 🔴 HIGH |
| | Top Edge: lateral_movement |
| | |
| | Influence per relation: |
| | lateral_movement +0.21031 |
| | has_vuln +0.14203 |
| | assumes_role +0.08821 |
| | accesses −0.03102 |
| | ``` |
| |
|
| | ## How to Load |
| |
|
| | ```python |
| | import torch, json |
| | from huggingface_hub import hf_hub_download |
| | |
| | weights_path = hf_hub_download(repo_id="adarsh-aur/Heterogenous_GNN", filename="model_final.pt") |
| | config_path = hf_hub_download(repo_id="adarsh-aur/Heterogenous_GNN", filename="config.json") |
| | |
| | with open(config_path) as f: |
| | cfg = json.load(f) |
| | |
| | model = HeteroRGCN( |
| | feat_dims = cfg["feat_dims"], |
| | node_types = cfg["node_types"], |
| | edge_types = [tuple(et) for et in cfg["edge_types"]], |
| | hidden_dim = cfg["hidden_dim"], |
| | out_dim = cfg["out_dim"], |
| | num_layers = cfg["num_layers"], |
| | dropout = cfg["dropout"], |
| | adapter_rank = cfg["adapter_rank"], |
| | predict_types = cfg["predict_types"], |
| | ) |
| | model.load_state_dict(torch.load(weights_path, map_location="cpu")) |
| | model.eval() |
| | ``` |
| |
|
| | ## Pipeline Position |
| |
|
| | | Stage | Component | |
| | |---|---| |
| | | Stage 5 | Graph Construction + Feature Fusion | |
| | | **Stage 6** | **Heterogeneous Structural GNN ← this model** | |
| | | Stage 7 | Temporal GNN | |
| | | Stage 8 | Risk Fusion MLP | |
| |
|
| | ## Training Details |
| |
|
| | | Property | Value | |
| | |---|---| |
| | | Optimizer | AdamW | |
| | | Scheduler | CosineAnnealingWarmRestarts | |
| | | Loss | BCEWithLogitsLoss + pos_weight (imbalance correction) | |
| | | Epochs | 150 | |
| | | Hardware | Kaggle T4 x2 | |
| | |