Publish 802 parameter graph compromise detector

README.md

@@ -0,0 +1,43 @@
+---
+license: apache-2.0
+task_categories:
+- tabular-classification
+tags:
+- graph-neural-network
+- cybersecurity
+- node-classification
+- pytorch
+---
+
+# MeshGraph GCN
+
+MeshGraph is a graph convolutional node classifier for detecting compromised assets
+inside a simulated enterprise network. Nodes carry host telemetry, while edges encode
+observed communication. Compromise begins at sparse seeds and propagates stochastically
+through the network.
+
+The GCN is compared with a logistic-regression baseline that sees identical host
+features but cannot use graph structure.
+
+This is a transductive benchmark: the graph and all node features are visible during
+training, while validation and test labels remain hidden.
+
+## Reproduce
+
+```powershell
+uv run python projects/meshgraph-gcn/generate_data.py
+uv run python projects/meshgraph-gcn/train.py
+```
+
+## Verified results
+
+The generated graph contains 600 nodes, 1,689 undirected communication edges, six
+subnets, and a 23.33% compromise rate. The final test contains 120 nodes:
+
+| Model | Parameters | Accuracy | ROC-AUC | Average precision | F1 |
+| --- | ---: | ---: | ---: | ---: | ---: |
+| Feature-only logistic regression | 9 fitted coefficients | 78.33% | 0.8362 | 0.6274 | 0.5938 |
+| Two-layer GCN | 802 | **81.67%** | **0.8564** | **0.6976** | **0.6333** |
+
+Both thresholds were selected independently on the same 120-node validation split.
+The GCN improved F1 by 3.96 points and average precision by 7.02 points.

evaluation.json

@@ -0,0 +1,45 @@
+{
+  "model": "MeshGraph GCN",
+  "parameters": 802,
+  "nodes": 600,
+  "edges": 1689,
+  "best_epoch": 26,
+  "best_validation_roc_auc": 0.9118788819875777,
+  "gcn_threshold": 0.6789770509686359,
+  "gcn_test": {
+    "accuracy": 0.8166666666666667,
+    "roc_auc": 0.8563664596273292,
+    "average_precision": 0.6975512110914732,
+    "precision": 0.59375,
+    "recall": 0.6785714285714286,
+    "f1": 0.6333333333333333,
+    "confusion_matrix": [
+      [
+        79,
+        13
+      ],
+      [
+        9,
+        19
+      ]
+    ]
+  },
+  "feature_only_logistic_test": {
+    "accuracy": 0.7833333333333333,
+    "roc_auc": 0.8361801242236025,
+    "average_precision": 0.6273562320739494,
+    "precision": 0.5277777777777778,
+    "recall": 0.6785714285714286,
+    "f1": 0.59375,
+    "confusion_matrix": [
+      [
+        75,
+        17
+      ],
+      [
+        9,
+        19
+      ]
+    ]
+  }
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3eec69401c9b743d6a0c72fa475710a847ba695ece1e31dc2517accd3a1492f7
+size 3512

preprocessing.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:87c61d8ec7ed2cd317b5becb844b767121937b1f35dae27f723155897b4d4f66
+size 568

source/generate_data.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import numpy as np
+from sklearn.model_selection import train_test_split
+
+PROJECT_DIR = Path(__file__).resolve().parent
+DATA_DIR = PROJECT_DIR / "data"
+
+
+def build_graph(nodes: int = 600, seed: int = 2033) -> dict[str, np.ndarray]:
+    rng = np.random.default_rng(seed)
+    subnets = np.repeat(np.arange(6), nodes // 6)
+    adjacency = np.zeros((nodes, nodes), dtype=np.float32)
+    for left in range(nodes):
+        same_subnet = subnets == subnets[left]
+        probabilities = np.where(same_subnet, 0.045, 0.0025)
+        links = rng.random(nodes) < probabilities
+        links[: left + 1] = False
+        adjacency[left, links] = 1
+    adjacency = np.maximum(adjacency, adjacency.T)
+
+    compromised = np.zeros(nodes, dtype=bool)
+    seeds = rng.choice(nodes, size=14, replace=False)
+    compromised[seeds] = True
+    for _ in range(4):
+        exposure = adjacency @ compromised.astype(np.float32)
+        infection_probability = 1 - np.exp(-0.22 * exposure)
+        new_infections = (rng.random(nodes) < infection_probability) & ~compromised
+        compromised |= new_infections
+
+    base = rng.normal(0, 1, (nodes, 8)).astype(np.float32)
+    labels = compromised.astype(np.int64)
+    signal = labels[:, None].astype(np.float32)
+    features = base.copy()
+    features[:, 0:1] += signal * rng.normal(1.0, 0.5, (nodes, 1))
+    features[:, 1:2] += signal * rng.normal(0.8, 0.6, (nodes, 1))
+    features[:, 2:3] += signal * rng.normal(0.7, 0.6, (nodes, 1))
+    features[:, 3:4] += signal * rng.normal(0.5, 0.7, (nodes, 1))
+    features[:, 4] += subnets * 0.12
+
+    indices = np.arange(nodes)
+    train, remainder = train_test_split(
+        indices,
+        test_size=0.40,
+        stratify=labels,
+        random_state=seed,
+    )
+    validation, test = train_test_split(
+        remainder,
+        test_size=0.50,
+        stratify=labels[remainder],
+        random_state=seed,
+    )
+    return {
+        "features": features,
+        "adjacency": adjacency,
+        "labels": labels,
+        "subnets": subnets.astype(np.int64),
+        "train_indices": train,
+        "validation_indices": validation,
+        "test_indices": test,
+        "seed_nodes": seeds.astype(np.int64),
+    }
+
+
+def main() -> None:
+    DATA_DIR.mkdir(parents=True, exist_ok=True)
+    graph = build_graph()
+    np.savez_compressed(DATA_DIR / "meshgraph.npz", **graph)
+    manifest = {
+        "nodes": len(graph["labels"]),
+        "edges": int(graph["adjacency"].sum() // 2),
+        "features": graph["features"].shape[1],
+        "subnets": len(np.unique(graph["subnets"])),
+        "compromised_rate": float(graph["labels"].mean()),
+        "train_nodes": len(graph["train_indices"]),
+        "validation_nodes": len(graph["validation_indices"]),
+        "test_nodes": len(graph["test_indices"]),
+        "path": "meshgraph.npz",
+    }
+    (DATA_DIR / "manifest.json").write_text(
+        json.dumps(manifest, indent=2),
+        encoding="utf-8",
+    )
+    print(json.dumps(manifest, indent=2))
+
+
+if __name__ == "__main__":
+    main()

source/model.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+class MeshGraphGCN(nn.Module):
+    def __init__(self, features: int = 8) -> None:
+        super().__init__()
+        self.input = nn.Linear(features, 32, bias=False)
+        self.hidden = nn.Linear(32, 16, bias=False)
+        self.output = nn.Linear(16, 2)
+
+    def forward(
+        self,
+        features: torch.Tensor,
+        normalized_adjacency: torch.Tensor,
+    ) -> torch.Tensor:
+        hidden = normalized_adjacency @ features
+        hidden = F.gelu(self.input(hidden))
+        hidden = F.dropout(hidden, p=0.15, training=self.training)
+        hidden = normalized_adjacency @ hidden
+        hidden = F.gelu(self.hidden(hidden))
+        return self.output(hidden)
+
+
+def normalize_adjacency(adjacency: torch.Tensor) -> torch.Tensor:
+    with_self_loops = adjacency + torch.eye(len(adjacency))
+    degree = with_self_loops.sum(dim=1).clamp(min=1)
+    inverse_sqrt = degree.pow(-0.5)
+    return inverse_sqrt[:, None] * with_self_loops * inverse_sqrt[None, :]
+
+
+def parameter_count(model: nn.Module) -> int:
+    return sum(parameter.numel() for parameter in model.parameters())

source/train.py

@@ -0,0 +1,191 @@
+from __future__ import annotations
+
+import json
+import random
+from pathlib import Path
+
+import numpy as np
+import torch
+import trackio
+from model import MeshGraphGCN, normalize_adjacency, parameter_count
+from safetensors.torch import save_file
+from sklearn.linear_model import LogisticRegression
+from sklearn.metrics import (
+    accuracy_score,
+    average_precision_score,
+    confusion_matrix,
+    f1_score,
+    precision_score,
+    recall_score,
+    roc_auc_score,
+)
+from torch.nn import functional as F
+
+PROJECT_DIR = Path(__file__).resolve().parent
+DATA_DIR = PROJECT_DIR / "data"
+ARTIFACT_DIR = PROJECT_DIR / "artifacts" / "meshgraph-gcn"
+
+
+def seed_everything(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+
+
+def metrics(labels: np.ndarray, scores: np.ndarray, threshold: float = 0.5) -> dict:
+    predictions = scores >= threshold
+    return {
+        "accuracy": float(accuracy_score(labels, predictions)),
+        "roc_auc": float(roc_auc_score(labels, scores)),
+        "average_precision": float(average_precision_score(labels, scores)),
+        "precision": float(precision_score(labels, predictions, zero_division=0)),
+        "recall": float(recall_score(labels, predictions, zero_division=0)),
+        "f1": float(f1_score(labels, predictions, zero_division=0)),
+        "confusion_matrix": confusion_matrix(labels, predictions).tolist(),
+    }
+
+
+def best_threshold(labels: np.ndarray, scores: np.ndarray) -> float:
+    candidates = np.quantile(scores, np.linspace(0.05, 0.95, 300))
+    return float(
+        max(
+            candidates,
+            key=lambda threshold: f1_score(
+                labels,
+                scores >= threshold,
+                zero_division=0,
+            ),
+        )
+    )
+
+
+def main() -> None:
+    seed_everything(2033)
+    graph = np.load(DATA_DIR / "meshgraph.npz")
+    features = graph["features"].astype(np.float32)
+    labels = graph["labels"].astype(np.int64)
+    train_indices = graph["train_indices"]
+    validation_indices = graph["validation_indices"]
+    test_indices = graph["test_indices"]
+    mean = features[train_indices].mean(axis=0)
+    scale = np.maximum(features[train_indices].std(axis=0), 1e-5)
+    features = (features - mean) / scale
+    feature_tensor = torch.from_numpy(features)
+    label_tensor = torch.from_numpy(labels)
+    adjacency = torch.from_numpy(graph["adjacency"].astype(np.float32))
+    normalized = normalize_adjacency(adjacency)
+
+    baseline = LogisticRegression(
+        class_weight="balanced",
+        max_iter=2000,
+        random_state=2033,
+    )
+    baseline.fit(features[train_indices], labels[train_indices])
+    baseline_validation = baseline.predict_proba(features[validation_indices])[:, 1]
+    baseline_threshold = best_threshold(
+        labels[validation_indices],
+        baseline_validation,
+    )
+    baseline_test = baseline.predict_proba(features[test_indices])[:, 1]
+
+    model = MeshGraphGCN(features=features.shape[1])
+    positive_weight = (labels[train_indices] == 0).sum() / max(
+        1, (labels[train_indices] == 1).sum()
+    )
+    class_weights = torch.tensor([1.0, positive_weight], dtype=torch.float32)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=0.01, weight_decay=0.002)
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=350)
+    best_validation_auc = -1.0
+    best_epoch = 0
+    best_state = None
+    trackio.init(
+        project="meshgraph-gcn",
+        name="two-layer-gcn-v1",
+        config={
+            "parameters": parameter_count(model),
+            "nodes": len(labels),
+            "edges": int(adjacency.sum().item() // 2),
+            "train_labels": len(train_indices),
+            "transductive": True,
+        },
+    )
+    for epoch in range(1, 351):
+        model.train()
+        logits = model(feature_tensor, normalized)
+        loss = F.cross_entropy(
+            logits[train_indices],
+            label_tensor[train_indices],
+            weight=class_weights,
+        )
+        optimizer.zero_grad(set_to_none=True)
+        loss.backward()
+        optimizer.step()
+        scheduler.step()
+        model.eval()
+        with torch.no_grad():
+            probabilities = model(feature_tensor, normalized).softmax(dim=1)[:, 1]
+        validation_auc = roc_auc_score(
+            labels[validation_indices],
+            probabilities[validation_indices].numpy(),
+        )
+        if validation_auc > best_validation_auc:
+            best_validation_auc = validation_auc
+            best_epoch = epoch
+            best_state = {
+                key: value.detach().cpu().clone()
+                for key, value in model.state_dict().items()
+            }
+        if epoch == 1 or epoch % 10 == 0:
+            trackio.log(
+                {
+                    "epoch": epoch,
+                    "train_loss": float(loss.detach()),
+                    "validation_roc_auc": validation_auc,
+                    "learning_rate": scheduler.get_last_lr()[0],
+                }
+            )
+    trackio.finish()
+    assert best_state is not None
+    model.load_state_dict(best_state)
+    model.eval()
+    with torch.no_grad():
+        gcn_scores = model(feature_tensor, normalized).softmax(dim=1)[:, 1].numpy()
+    gcn_threshold = best_threshold(
+        labels[validation_indices],
+        gcn_scores[validation_indices],
+    )
+    results = {
+        "model": "MeshGraph GCN",
+        "parameters": parameter_count(model),
+        "nodes": len(labels),
+        "edges": int(adjacency.sum().item() // 2),
+        "best_epoch": best_epoch,
+        "best_validation_roc_auc": best_validation_auc,
+        "gcn_threshold": gcn_threshold,
+        "gcn_test": metrics(
+            labels[test_indices],
+            gcn_scores[test_indices],
+            gcn_threshold,
+        ),
+        "feature_only_logistic_test": metrics(
+            labels[test_indices],
+            baseline_test,
+            baseline_threshold,
+        ),
+    }
+    ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
+    save_file(model.state_dict(), ARTIFACT_DIR / "model.safetensors")
+    np.savez(
+        ARTIFACT_DIR / "preprocessing.npz",
+        mean=mean,
+        scale=scale,
+    )
+    (ARTIFACT_DIR / "evaluation.json").write_text(
+        json.dumps(results, indent=2),
+        encoding="utf-8",
+    )
+    print(json.dumps(results, indent=2))
+
+
+if __name__ == "__main__":
+    main()