Graph Neural Network Models

Spec-layer graph neural-network model definitions.

The layer-level message-passing math lives in NN.Spec.Layers.Gnn, in particular Spec.GCNLayerSpec, Spec.gcnLayerSpec, and Spec.gcnLayerBackwardSpec. This module wires those layers into a two-layer GCN with graph-level mean pooling and records the end-to-end shapes.

Reference (GCN):

• Kipf and Welling, "Semi-Supervised Classification with Graph Convolutional Networks" (2017): https://arxiv.org/abs/1609.02907

PyTorch ecosystem analogies:

• torch_geometric.nn.GCNConv for the layer-level GCN operator,
• global mean pooling as in torch_geometric.nn.global_mean_pool.

Implementation status

No API builder implements this model (nn.* has no graph layers), no runtime code imports it, and no theorem relates it to anything else. It is a standalone reference specification.

2-layer GCN with gradients

Forward is:

1. Z₁ = GCN₁(X) (linear message passing)
2. H₁ = ReLU(Z₁)
3. H₂ = GCN₂(H₁)
4. y = mean_nodes(H₂) (graph-level readout)

Diagram (single graph, no batching):

X : (n × inDim)
   └─ GCN₁ ─→ Z₁ : (n × hidDim) ─ ReLU ─→ H₁ : (n × hidDim)
                                 └─ GCN₂ ─→ H₂ : (n × outDim)
                                              └─ mean over nodes ─→ y : (outDim)

Backward mirrors this structure:

• "mean nodes" broadcasts the outDim gradient back to n×outDim and scales by 1/n,
• each GCN layer uses the matrix-calculus rules in Spec.gcnLayerBackwardSpec,
• ReLU gates gradients by ReLU'.

structure Models.GCN2Spec (n inDim hidDim outDim : ℕ) (α : Type) [TorchLean.Storage α] : Type

A 2-layer GCN "model spec" for a fixed graph with n nodes.

GCNLayerSpec packages the per-layer parameters (including the adjacency/normalization choice), so the model here is just two such layers composed with a nonlinearity and a readout.

• inputLayer : Spec.GCNLayerSpec n inDim hidDim α

  First GCN layer: inDim → hidDim.

• outputLayer : Spec.GCNLayerSpec n hidDim outDim α

  Second GCN layer: hidDim → outDim.

def Models.GCN2Spec.forward {α : Type} [TorchLean.Storage α] [Context α] {n inDim hidDim outDim : ℕ} (m : GCN2Spec n inDim hidDim outDim α) (x : TorchLean.Tensor α [n, inDim]) (h_n : n > 0) : TorchLean.Tensor α [outDim]

Forward pass for the 2-layer GCN with a graph-level mean pooling readout.

Input:

• x : n × inDim node features.

Output:

• y : outDim graph embedding produced by averaging node embeddings.

The h_n : n > 0 assumption is only used to make the mean pooling well-defined (division by n).

structure Models.GCN2Grads (n inDim hidDim outDim : ℕ) (α : Type) [TorchLean.Storage α] : Type

Gradients for both layers of GCN2Spec.

• inputLayer : Spec.GCNLayerParameterGradients n inDim hidDim α

  Gradients for the inDim → hidDim input GCN layer.

• outputLayer : Spec.GCNLayerParameterGradients n hidDim outDim α

  Gradients for the hidDim → outDim output GCN layer.

def Models.GCN2Spec.backward {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {n inDim hidDim outDim : ℕ} (m : GCN2Spec n inDim hidDim outDim α) (x : TorchLean.Tensor α [n, inDim]) (gradOutput : TorchLean.Tensor α [outDim]) (h_n : n > 0) : GCN2Grads n inDim hidDim outDim α × TorchLean.Tensor α [n, inDim]

Backward/VJP for GCN2Spec.forward.

This is written in the same "spec style" as the rest of TorchLean:

• recompute small intermediates instead of depending on runtime caches,
• apply VJPs in reverse order,
• keep shapes explicit.

PyTorch analogy: this corresponds to what autograd would do for a graph built from GCNConv → ReLU → GCNConv → global_mean_pool, but expressed as a pure function.