adarsh-aur/Heterogenous_GNN

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

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