Graph Neural Networks 🕸️

Many real-world problems are naturally graphs:

• UPI fraud: accounts ↔ transactions
• Drug discovery: atoms ↔ chemical bonds
• Recommendations: users ↔ products they bought
• Knowledge graphs: entities ↔ relationships
• Indian Railways: stations ↔ train routes (delay propagation)

The core operation in a GNN is message passing: each node updates its representation by aggregating information from its neighbours. Stack k layers, and each node "sees" k hops away.

!pip install -q torch-geometric

Build the transaction graph:

import torch
from torch_geometric.data import Data

# Each account is a node with features (age_days, kyc_level, avg_balance, ...)
node_features = torch.tensor([
    [180, 2, 12000.0, 23, 0],   # account 0: 180 days old, KYC L2, ...
    [12,  1, 800.0,   45, 1],   # account 1: new, low KYC, suspicious
    # ... thousands of accounts
], dtype=torch.float)

# Edges: each transaction is a directed edge from sender to receiver
edge_index = torch.tensor([
    [0, 1, 1, 2, 3],   # source accounts
    [1, 2, 3, 4, 4],   # destination accounts
], dtype=torch.long)

# Edge features: amount, timestamp, channel
edge_attr = torch.tensor([
    [5000.0, 1700000000, 0],
    [4900.0, 1700000050, 0],
    # ...
], dtype=torch.float)

# Labels: 0 = clean, 1 = fraud (only on labelled accounts)
y = torch.tensor([0, 1, 1, 0, 0], dtype=torch.long)

graph = Data(x=node_features, edge_index=edge_index,
             edge_attr=edge_attr, y=y)

Define the GraphSAGE classifier:

import torch.nn.functional as F
from torch_geometric.nn import SAGEConv

class FraudGNN(torch.nn.Module):
    def __init__(self, in_dim, hidden_dim, out_dim):
        super().__init__()
        self.sage1 = SAGEConv(in_dim, hidden_dim, aggr="mean")
        self.sage2 = SAGEConv(hidden_dim, hidden_dim, aggr="mean")
        self.classifier = torch.nn.Linear(hidden_dim, out_dim)

    def forward(self, x, edge_index):
        # 1-hop neighbourhood aggregation
        h = self.sage1(x, edge_index)
        h = F.relu(h)
        h = F.dropout(h, p=0.3, training=self.training)
        # 2-hop: each node now sees neighbours-of-neighbours
        h = self.sage2(h, edge_index)
        h = F.relu(h)
        return self.classifier(h)  # logits per node

Real fraud graphs have millions of nodes — too big to fit one batch. NeighborLoader samples a fixed-size neighbourhood per labelled node:

from torch_geometric.loader import NeighborLoader

train_loader = NeighborLoader(
    graph,
    num_neighbors=[20, 10],   # 20 1-hop neighbours, 10 2-hop neighbours
    batch_size=256,
    input_nodes=train_mask,
)

model = FraudGNN(in_dim=5, hidden_dim=64, out_dim=2).cuda()
opt = torch.optim.Adam(model.parameters(), lr=1e-3)

# Class weighting — fraud is rare (~0.5% of accounts)
class_weights = torch.tensor([1.0, 50.0]).cuda()

for epoch in range(20):
    model.train()
    for batch in train_loader:
        batch = batch.cuda()
        out = model(batch.x, batch.edge_index)
        # Predictions only for the seed nodes in the batch
        loss = F.cross_entropy(
            out[:batch.batch_size],
            batch.y[:batch.batch_size],
            weight=class_weights,
        )
        opt.zero_grad(); loss.backward(); opt.step()
    print(f"Epoch {epoch}  loss={loss.item():.4f}")

• Inference latency: Each new transaction adds an edge. To score it, gather both accounts' k-hop neighbourhoods, run the GNN. With k=2 and degree cap 20, this is ~30ms on CPU.
• Online graph updates: Use a graph store like Neo4j or DGraph that handles streaming edge inserts. Embed-on-the-fly for new nodes via GraphSAGE's inductive nature.
• Explainability: Use GNNExplainer from PyG to highlight which neighbours drove the fraud prediction — auditors and regulators need this.
• Cold start: Brand-new accounts have no neighbours. Fall back to feature-only logistic regression for the first 7 days.