Graph Neural Networks: Architectures, Training, and Production Deployment

Beyond Grids and Sequences

Deep learning conquered images with convolutions and language with transformers. Both domains assume a regular structure — pixels in a grid, tokens in a sequence. But the world's most interesting data lives in graphs: social networks, molecular structures, citation networks, knowledge graphs, and supply chains. These have no fixed grid, no sequential order. Each graph has a different number of nodes with a different arrangement of edges.

Graph Neural Networks (GNNs)¹ extend deep learning to this irregular domain. Instead of sliding a kernel over a grid, GNNs propagate information along edges, letting each node aggregate features from its neighbours. The result is a representation that captures both the node's own properties and the structure of its local neighbourhood — exactly what you need for tasks like predicting protein interactions, ranking documents in a knowledge graph, or detecting fraud in a transaction network.

This article surveys the major GNN architectures, shows you how to train them with PyTorch Geometric, and covers what it takes to deploy them to production.

The Message-Passing Paradigm

Every GNN is built on the same core operation: message passing. For each node, the GNN collects messages from its neighbours, aggregates them (sum, mean, or attention-weighted), and updates the node's representation. After one round, each node knows about its direct neighbours. After k rounds, each node knows about nodes up to k hops away.

h_v^(0) = x_v                          (initial node features)

h_v^(k) = UPDATE(
    h_v^(k-1),
    AGGREGATE({ h_u^(k-1) for every neighbour u of v })
)

This is deceptively simple. The choice of aggregate and update functions determines the entire GNN's behaviour, and different architectures make different trade-offs between expressiveness, computational cost, and over-smoothing.

GNN Architectures Compared

Architecture	Aggregation	Key Innovation	Best For	Limitations
GCN (Kipf & Welling, 2017)	Normalised mean	Symmetric normalisation of adjacency matrix	Citation networks, node classification	Treats all neighbours equally
GraphSAGE (Hamilton et al., 2017)	Mean / LSTM / Pool	Inductive — learns for unseen nodes; supports sampling	Large graphs, evolving graphs	Higher memory than GCN with LSTM aggregator
GAT (Veličković et al., 2018)	Attention-weighted sum	Self-attention over neighbours — learns which neighbours matter	Heterogeneous graphs, relationship-rich domains	O(E) attention computation
GIN (Xu et al., 2019)	Sum + MLP	Maximally expressive under WL-test	Graph classification, isomorphism detection	Deeper networks still over-smooth
Graph Transformer (Dwivedi & Bresson, 2021)	Full attention	Node positional encoding + transformer attention	Large graphs with long-range dependencies	O(V²) attention — does not scale beyond ~5k nodes

GCN is where most people start. It is simple, fast, and works well for homophilic graphs — where connected nodes tend to be similar (citation networks, social networks). GAT is the go-to when edges carry different meanings: if a "friend" edge should matter more than a "follower" edge, attention weights learn that automatically. GraphSAGE is essential for production systems because it supports inductive inference — you can train on one graph and run inference on a completely unseen graph. Graph Transformers trade quadratic memory for the ability to capture long-range dependencies that shallow message passing cannot reach.

Training Pipeline with PyTorch Geometric

PyTorch Geometric (PyG) is the de facto standard for GNN training. Here is a complete training loop for node classification on the Cora citation network:

import torch
import torch.nn.functional as F
from torch_geometric.datasets import Planetoid
from torch_geometric.nn import GCNConv

# 1. Load data
dataset = Planetoid(root="/tmp/cora", name="Cora")
data = dataset[0]

# 2. Define model
class GCN(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = GCNConv(in_channels, hidden_channels)
        self.conv2 = GCNConv(hidden_channels, out_channels)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, training=self.training, p=0.5)
        x = self.conv2(x, edge_index)
        return F.log_softmax(x, dim=1)

model = GCN(
    in_channels=dataset.num_features,
    hidden_channels=16,
    out_channels=dataset.num_classes,
)
optimiser = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)

# 3. Training loop
def train():
    model.train()
    optimiser.zero_grad()
    out = model(data.x, data.edge_index)
    loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimiser.step()
    return float(loss)

# 4. Evaluation
@torch.no_grad()
def test():
    model.eval()
    out = model(data.x, data.edge_index)
    pred = out.argmax(dim=1)
    accs = []
    for mask in [data.train_mask, data.val_mask, data.test_mask]:
        acc = (pred[mask] == data.y[mask]).sum().item() / mask.sum().item()
        accs.append(acc)
    return accs

for epoch in range(200):
    loss = train()
    train_acc, val_acc, test_acc = test()
    if epoch % 20 == 0:
        print(f"Epoch {epoch:3d} | Loss: {loss:.4f} | "
              f"Train: {train_acc:.3f} | Val: {val_acc:.3f} | Test: {test_acc:.3f}")

At 200 epochs on Cora (2,708 nodes, 5,429 edges), this GCN reaches roughly 81% test accuracy. The same architecture scaled to 100k nodes with GraphSAGE's neighbour sampling runs in minutes on a single GPU.

Choosing a GNN Framework

The examples above use PyTorch Geometric (PyG), but it is not the only option. The framework you choose affects scalability, hardware support, and production ergonomics:

Framework	Backend	API Style	Scalability	Best For
PyTorch Geometric	PyTorch	Modular layers	Moderate (custom samplers)	Rapid prototyping, research, heterogeneous GNNs
Deep Graph Library (DGL)	PyTorch / TensorFlow	Message-passing primitives	High (built-in sampling, distributed)	Large-scale graphs, production pipelines
Jraph	JAX	Functional transformations	High (TPU-native)	TPU training, JAX ecosystem, differentiable programs
TensorFlow GNN	TensorFlow	Declarative config	High (TFX integration)	TF-native stacks, TF-serving deployment

PyG offers the richest collection of pre-built convolutional layers (GCN, GAT, GIN, RGCN, and dozens more) and the largest community, making it the best starting point for most projects. DGL pulls ahead when graph size demands advanced neighbour sampling or multi-GPU training — its dgl.dataloading.NeighborSampler pipeline is battle-tested on graphs exceeding 100 million nodes. Jraph is the choice for teams already on JAX who need end-to-end differentiability and TPU acceleration; the trade-off is a smaller ecosystem and steeper learning curve. TensorFlow GNN makes sense only when the broader stack is TensorFlow-native, because its graph-model definition uses protocol buffers and a declarative config language that diverges from the intuitive Python-first APIs of PyG and DGL.

GNNs for Knowledge Graphs

GNNs have a natural overlap with knowledge graphs. A knowledge graph is a heterogeneous directed graph — nodes have types (Person, Company, Patent) and edges have labelled roles (WORKS_AT, FILES, CITES). Standard GNNs assume a single edge type with undirected semantics, so applying them to knowledge graphs requires extensions:

• R-GCN (Schlichtkrull et al., 2018) maintains a separate transformation matrix per relation type. This works well for small-to-medium knowledge graphs but scales poorly — O(R × d²) parameters where R is the number of relation types and d is the hidden dimension.
• CompGCN (Vashishth et al., 2020) composes edge embeddings with node embeddings during message passing, reducing the parameter count. It handles 10× more relation types than R-GCN for the same budget.
• Link prediction in knowledge graphs uses a scoring function (DistMult, ComplEx, ConvE) on the node embeddings produced by the GNN. The model learns to score existing triples (h, r, t) higher than corrupted ones — the same loss structure used by TransE and its successors.

The result is a model that can predict missing links in your knowledge graph: "Which companies are likely to form partnerships?" or "What technologies might a given organisation develop next?"

Heterogeneous GNNs in Practice: R-GCN for Link Prediction

To ground the theory, here is a complete R-GCN training loop for link prediction on a small synthetic knowledge graph. The model learns to score triples (head, relation, tail) — the fundamental unit of any knowledge graph — using DistMult as the scoring function:

import torch
import torch.nn.functional as F
from torch_geometric.nn import RGCNConv, DistMult

class RGCNLinkPredictor(torch.nn.Module):
    def __init__(self, num_nodes, num_rels, hidden_dim=64):
        super().__init__()
        self.node_emb = torch.nn.Embedding(num_nodes, hidden_dim)
        self.conv1 = RGCNConv(hidden_dim, hidden_dim, num_rels)
        self.conv2 = RGCNConv(hidden_dim, hidden_dim, num_rels)
        self.scorer = DistMult(hidden_dim, num_rels)

    def forward(self, edge_index, edge_type):
        x = self.node_emb.weight
        x = self.conv1(x, edge_index, edge_type).relu()
        x = F.dropout(x, p=0.2, training=self.training)
        x = self.conv2(x, edge_index, edge_type)
        return x  # node embeddings

    def score_triples(self, head, rel, tail, emb=None):
        if emb is None:
            emb = self.forward(None, None)
        return self.scorer(emb[head], rel, emb[tail])

# Training uses 1‑vs‑all negative sampling — corrupt either head or tail
# and minimise the margin ranking loss between positive and negative scores.

The key design choice is the scoring function. DistMult is parameter-efficient and works well for symmetric relations. For asymmetric relations (e.g., EMPLOYS), ComplEx's complex-valued embeddings capture directionality better. In production, benchmark all three — DistMult, ComplEx, and ConvE — on a held-out validation set before committing.

Training R-GCN at scale requires the same neighbourhood sampling techniques used by GraphSAGE. PyG's RGCNConv supports GraphSAINT-style sampling natively, which reduces the effective computation graph from the full KG to a sampled subgraph per mini-batch. On a knowledge graph with 500k entities and 100 relation types, sampling 2‑hop neighbourhoods of 256 nodes per batch keeps GPU memory under 8 GB.

Choosing a GNN Framework

PyTorch Geometric (PyG) and Deep Graph Library (DGL) dominate the GNN framework landscape. The choice between them depends on your deployment constraints:

Criterion	PyG	DGL
API style	PyTorch-native (nn.Module subclassing)	Computation graph abstraction (graph.update_all())
Heterogeneous graphs	RGCNConv, HeteroConv, HANConv — mature	RelGraphConv, HeteroGraphConv — mature
Sampling	NeighborSampler, GraphSAINT, Shadow	MultiLayerNeighborSampler, ClusterGCN, SAINTSampler
Scalability	Excellent; used in production at Twitter/Meta	Excellent; used in production at Amazon/Apple
ONNX export	Supported via torch.onnx.export — some custom ops may fail	Native ONNX support with DGL's dgl.batched_graph_to_onnx
Community & ecosystem	Larger research community; most SOTA papers release PyG code first	Strong industrial adoption; better integration with AWS SageMaker
Learning curve	Shallow — familiar PyTorch patterns	Steeper — requires understanding DGL's message-passing primitives

In practice, both frameworks converge on the same feature set. PyG is the better choice for teams already embedded in the PyTorch ecosystem. DGL pulls ahead when you need deep integration with AWS infrastructure or when ONNX export is a hard requirement.

The Graph Transformer Frontier

The architectures table above position Graph Transformers as a niche tool due to O(V²) attention. However, two developments have broadened their applicability in 2025–2026:

• Linear attention variants (Performer, BigBird) reduce complexity to O(V), making Graph Transformers viable for graphs up to ~50k nodes. GPS (General, Powerful, Scalable) Graph Transformer combines message-passing layers with efficient linear attention, achieving state-of-the-art results on the long-range graph benchmark (LRGB) while remaining computationally tractable.
• Graph Transformers for heterophilic graphs. Standard GNNs assume homophily — connected nodes are similar. On heterophilic graphs (e.g., molecular structures where atom types are distributed by bonding constraints), message passing can actually harm performance by smoothing node representations across dissimilar neighbours. Graph Transformers, which attend to all nodes (or a learned subset) rather than only neighbours, consistently outperform MP-GNNs on heterophilic benchmarks.

Practical guidance: Start with GCN or GAT on homophilic graphs up to 1M nodes. Switch to GraphSAGE for inductive or streaming settings. Reserve Graph Transformers for graphs with long-range dependencies (molecules, code ASTs) or when node homophily is below 0.3. For everything in between, GPS Graph Transformer with linear attention offers a robust default.

Production Deployment Considerations

Deploying GNNs to production is harder than training them. Four challenges dominate:

Challenge	Problem	Mitigation
Inference on evolving graphs	Nodes and edges are added constantly; re-training is expensive	GraphSAGE inductive inference; streaming feature updates
Neighbourhood explosion	Full-batch adjacency on a 10M-node graph does not fit in GPU memory	Neighbour sampling (GraphSAINT, ClusterGCN); distributed mini-batch training
Feature staleness	Node features drift over time, degrading prediction quality	Online feature stores (Feast, Tecton) with TTL-based recomputation
Explainability	Regulators require reasons for edge-level predictions	GNNExplainer, Integrated Gradients on graph structure; subgraph-level explanations

Neighbour sampling is the single most important optimisation for production GNNs. Instead of materialising the full computation graph, each mini-batch samples a fixed-size neighbourhood per node — typically 10–25 neighbours at 2–3 hops. Tools like PyG's NeighborSampler and DGL's MultiLayerNeighborSampler handle this natively. For a graph of 10M nodes with average degree 50, sampling reduces per-batch memory from terabytes to megabytes.

For production inference, frameworks like TorchServe and NVIDIA Triton serve PyG models with GPU batching. A well-optimised GraphSAGE model with neighbour sampling can achieve sub-100ms p99 latency on graphs of 1M+ nodes with a single V100 GPU.

From Training to Inference: Exporting Your GNN

The training loop above produces a model in memory. For production, you need a serialised artefact that can be versioned and served independently of the training code. TorchScript traces the model's forward pass into a self-contained graph that runs without Python dependencies:

import torch

# Trace the trained GCN into TorchScript
model.eval()
traced = torch.jit.trace(model, (data.x, data.edge_index))
traced.save("gnn_cora.pt")

# Load and run inference — no PyG import needed
loaded = torch.jit.load("gnn_cora.pt")
loaded.eval()
with torch.no_grad():
    logits = loaded(data.x, data.edge_index)
    predictions = logits.argmax(dim=1)
    print(f"Predicted {predictions.sum():,} nodes")

For GPU serving at production scale, export to ONNX and serve via NVIDIA Triton or ONNX Runtime — both support dynamic batching and model versioning. TorchServe handles CPU deployments with integrated metrics and A/B testing. The neighbour sampling strategies discussed earlier combine naturally with these serving frameworks: GraphSAGE models exported to ONNX and served on Triton achieve sub-100ms p99 latency on graphs exceeding one million nodes.

Further Reading

• Graph Theory for Software Engineers — The algorithmic foundations that GNNs build on, from BFS to PageRank
• Graph Algorithms in Production with Neo4j GDS — Classical graph algorithms for feature engineering before GNN training
• Knowledge Graph Construction from Unstructured Data — How to build the graphs that GNNs can learn from
• Community Detection Algorithms — Graph partitioning techniques used in GNN training and GraphRAG pipelines
• Introduction to GraphRAG — Graph-augmented retrieval where GNNs can improve entity ranking
• Ontology in Graph Databases — Schema design principles that make GNN features meaningful
• Graph-Powered OSINT — Applied entity resolution and relationship inference using graph analysis
• Knowledge Graphs as Context Layer for AI Agents — How GNN-generated embeddings feed agent memory and retrieval pipelines
• Community Detection Algorithms — Classical partitioning methods that complement GNN-based clustering
• PyTorch Geometric documentation — pytorch-geometric.readthedocs.io
• DGL documentation — dgl.ai

Footnotes

1. The term "Graph Neural Network" was first coined by Scarselli et al. in 2005, but the field only gained wide adoption after Kipf & Welling's GCN paper in 2017 and the release of PyTorch Geometric in 2019. ↩