Hugging Face's logo

adarsh-aur
/
Heterogenous_GNN

Graph Machine Learning
English
graph-neural-network
rgcn
heterogeneous-graph
cybersecurity
anomaly-detection
pytorch-geometric
multi-cloud
threat-detection
Model card Files
xet
Community
Heterogenous_GNN
File size: 2,925 Bytes
96f66b6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
 
---
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 |