Topic 20: Graph Neural Networks (GNNs) - First Principles¶

Goal: Implement Graph Laplacian, GCN, and GAT from scratch using pure NumPy. Prerequisites: Basic linear algebra and neural networks. Reference: theory.md

In [1]:

import random

import matplotlib.pyplot as plt
import numpy as np

%load_ext autoreload
%autoreload 2

SEED = 42
random.seed(SEED)
rng = np.random.default_rng(SEED)

plt.rcParams["figure.dpi"] = 90

import random import matplotlib.pyplot as plt import numpy as np %load_ext autoreload %autoreload 2 SEED = 42 random.seed(SEED) rng = np.random.default_rng(SEED) plt.rcParams["figure.dpi"] = 90

1. Problem Setup — WHY¶

We want to predict properties of nodes in a graph. For this implementation, we will use the famous Zachary's Karate Club graph. The club split into two factions, and we want to predict which faction each member joined based only on the graph topology.

In [2]:

# Build Zachary's Karate Club graph using pure NumPy
N = 34
edges = [
    (0,1), (0,2), (0,3), (0,4), (0,5), (0,6), (0,7), (0,8), (0,10), (0,11), (0,12), (0,13), (0,17), (0,19), (0,21), (0,31),
    (1,2), (1,3), (1,7), (1,13), (1,17), (1,19), (1,21), (1,30), (2,3), (2,7), (2,8), (2,9), (2,13), (2,27), (2,28), (2,32),
    (3,7), (3,12), (3,13), (4,6), (4,10), (5,6), (5,10), (5,16), (6,16), (8,30), (8,32), (8,33), (9,33), (13,33),
    (14,32), (14,33), (15,32), (15,33), (18,32), (18,33), (19,33), (20,32), (20,33), (22,32), (22,33), (23,25), (23,27),
    (23,29), (23,32), (23,33), (24,25), (24,27), (24,31), (25,31), (26,29), (26,33), (27,33), (28,31), (28,33), (29,32),
    (29,33), (30,32), (30,33), (31,32), (31,33), (32,33)
]

A = np.zeros((N, N))
for u, v in edges:
    A[u, v] = 1.0
    A[v, u] = 1.0

# Ground truth club split (0 for 'Mr. Hi', 1 for 'Officer')
labels = np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 1, 0, 0, 1, 0, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1])

# One-hot node identity features
X = np.eye(N)

print(f"Number of nodes: {N}")
print(f"Adjacency matrix shape: {A.shape}")
print(f"Feature matrix shape: {X.shape}")

# Build Zachary's Karate Club graph using pure NumPy N = 34 edges = [ (0,1), (0,2), (0,3), (0,4), (0,5), (0,6), (0,7), (0,8), (0,10), (0,11), (0,12), (0,13), (0,17), (0,19), (0,21), (0,31), (1,2), (1,3), (1,7), (1,13), (1,17), (1,19), (1,21), (1,30), (2,3), (2,7), (2,8), (2,9), (2,13), (2,27), (2,28), (2,32), (3,7), (3,12), (3,13), (4,6), (4,10), (5,6), (5,10), (5,16), (6,16), (8,30), (8,32), (8,33), (9,33), (13,33), (14,32), (14,33), (15,32), (15,33), (18,32), (18,33), (19,33), (20,32), (20,33), (22,32), (22,33), (23,25), (23,27), (23,29), (23,32), (23,33), (24,25), (24,27), (24,31), (25,31), (26,29), (26,33), (27,33), (28,31), (28,33), (29,32), (29,33), (30,32), (30,33), (31,32), (31,33), (32,33) ] A = np.zeros((N, N)) for u, v in edges: A[u, v] = 1.0 A[v, u] = 1.0 # Ground truth club split (0 for 'Mr. Hi', 1 for 'Officer') labels = np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 1, 0, 0, 1, 0, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]) # One-hot node identity features X = np.eye(N) print(f"Number of nodes: {N}") print(f"Adjacency matrix shape: {A.shape}") print(f"Feature matrix shape: {X.shape}")

Number of nodes: 34
Adjacency matrix shape: (34, 34)
Feature matrix shape: (34, 34)

2. Mathematical Core — WHAT¶

Key equations:

1. Unnormalized Laplacian: $L = D - A$
2. Normalized Adjacency (GCN): $\hat{A} = \tilde{D}^{-1/2} \tilde{A} \tilde{D}^{-1/2}$ where $\tilde{A} = A + I$
3. GCN Layer: $H^{(l+1)} = \sigma(\hat{A} H^{(l)} W^{(l)})$
4. GAT Attention: $\alpha_{ij} = \text{softmax}_j(\text{LeakyReLU}(a^T [Wh_i || Wh_j]))$

3. Solution Method — HOW¶

We will build the components bottom-up:

1. Utilities for Laplacian and GCN normalization.
2. A single GCN Layer.
3. A single GAT Layer.
4. A full network (2 layers).

4. Implementation — BUILD¶

In [3]:

def get_laplacian(A):
    D = np.diag(np.sum(A, axis=1))
    L = D - A
    return L

L = get_laplacian(A)
eigenvalues = np.linalg.eigvals(L)
print(f"Smallest eigenvalues: {np.sort(eigenvalues)[:5]}")
assert np.all(eigenvalues > -1e-10), "Laplacian should be Positive Semi-Definite"

def get_laplacian(A): D = np.diag(np.sum(A, axis=1)) L = D - A return L L = get_laplacian(A) eigenvalues = np.linalg.eigvals(L) print(f"Smallest eigenvalues: {np.sort(eigenvalues)[:5]}") assert np.all(eigenvalues > -1e-10), "Laplacian should be Positive Semi-Definite"

Smallest eigenvalues: [2.56113228e-16 4.68525227e-01 9.09247664e-01 1.12501072e+00
 1.25940411e+00]

In [4]:

def gcn_normalize(A):
    A_tilde = A + np.eye(A.shape[0])
    D_tilde_diag = np.sum(A_tilde, axis=1)
    # D^{-1/2}
    D_tilde_inv_sqrt = np.diag(np.power(D_tilde_diag, -0.5))
    # D^{-1/2} * A * D^{-1/2}
    A_norm = D_tilde_inv_sqrt @ A_tilde @ D_tilde_inv_sqrt
    return A_norm

A_norm = gcn_normalize(A)
print(f"A_norm shape: {A_norm.shape}")

def gcn_normalize(A): A_tilde = A + np.eye(A.shape[0]) D_tilde_diag = np.sum(A_tilde, axis=1) # D^{-1/2} D_tilde_inv_sqrt = np.diag(np.power(D_tilde_diag, -0.5)) # D^{-1/2} * A * D^{-1/2} A_norm = D_tilde_inv_sqrt @ A_tilde @ D_tilde_inv_sqrt return A_norm A_norm = gcn_normalize(A) print(f"A_norm shape: {A_norm.shape}")

A_norm shape: (34, 34)

In [5]:

def relu(x):
    return np.maximum(0, x)

def softmax(x, axis=-1):
    e_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
    return e_x / e_x.sum(axis=axis, keepdims=True)

class GCNLayer:
    def __init__(self, in_features, out_features, rng):
        # Xavier initialization
        limit = np.sqrt(6 / (in_features + out_features))
        self.W = rng.uniform(-limit, limit, size=(in_features, out_features))
        
    def forward(self, H, A_norm):
        # H: (N, in_features)
        # A_norm: (N, N)
        # W: (in_features, out_features)
        # H @ W -> (N, out_features)
        # A_norm @ (H @ W) -> (N, out_features)
        return A_norm @ H @ self.W

def relu(x): return np.maximum(0, x) def softmax(x, axis=-1): e_x = np.exp(x - np.max(x, axis=axis, keepdims=True)) return e_x / e_x.sum(axis=axis, keepdims=True) class GCNLayer: def __init__(self, in_features, out_features, rng): # Xavier initialization limit = np.sqrt(6 / (in_features + out_features)) self.W = rng.uniform(-limit, limit, size=(in_features, out_features)) def forward(self, H, A_norm): # H: (N, in_features) # A_norm: (N, N) # W: (in_features, out_features) # H @ W -> (N, out_features) # A_norm @ (H @ W) -> (N, out_features) return A_norm @ H @ self.W

In [6]:

class GATLayer:
    def __init__(self, in_features, out_features, rng, alpha=0.2):
        limit = np.sqrt(6 / (in_features + out_features))
        self.W = rng.uniform(-limit, limit, size=(in_features, out_features))
        self.a = rng.uniform(-limit, limit, size=(2 * out_features, 1))
        self.alpha = alpha # LeakyReLU slope
        
    def forward(self, H, A):
        N = H.shape[0]
        # Linear transformation
        WH = H @ self.W  # (N, out_features)
        
        # Prepare concatenations for all pairs
        # WH_i is repeated N times, WH_j is tiled N times
        WH_i = np.repeat(WH, N, axis=0).reshape(N, N, -1)
        WH_j = np.tile(WH, (N, 1)).reshape(N, N, -1)
        
        # Concatenate WH_i and WH_j
        WH_concat = np.concatenate([WH_i, WH_j], axis=-1)  # (N, N, 2*out_features)
        
        # Compute attention scores
        e = (WH_concat @ self.a).squeeze(-1)  # (N, N)
        
        # LeakyReLU
        e = np.where(e > 0, e, self.alpha * e)
        
        # Mask out non-neighbors (use large negative value for softmax)
        # Include self-loops
        A_tilde = A + np.eye(N)
        mask = np.where(A_tilde > 0, 0, -1e9)
        e = e + mask
        
        # Softmax
        attention = softmax(e, axis=1)  # (N, N)
        
        # Aggregate
        H_out = attention @ WH
        return H_out, attention

class GATLayer: def __init__(self, in_features, out_features, rng, alpha=0.2): limit = np.sqrt(6 / (in_features + out_features)) self.W = rng.uniform(-limit, limit, size=(in_features, out_features)) self.a = rng.uniform(-limit, limit, size=(2 * out_features, 1)) self.alpha = alpha # LeakyReLU slope def forward(self, H, A): N = H.shape[0] # Linear transformation WH = H @ self.W # (N, out_features) # Prepare concatenations for all pairs # WH_i is repeated N times, WH_j is tiled N times WH_i = np.repeat(WH, N, axis=0).reshape(N, N, -1) WH_j = np.tile(WH, (N, 1)).reshape(N, N, -1) # Concatenate WH_i and WH_j WH_concat = np.concatenate([WH_i, WH_j], axis=-1) # (N, N, 2*out_features) # Compute attention scores e = (WH_concat @ self.a).squeeze(-1) # (N, N) # LeakyReLU e = np.where(e > 0, e, self.alpha * e) # Mask out non-neighbors (use large negative value for softmax) # Include self-loops A_tilde = A + np.eye(N) mask = np.where(A_tilde > 0, 0, -1e9) e = e + mask # Softmax attention = softmax(e, axis=1) # (N, N) # Aggregate H_out = attention @ WH return H_out, attention

5. Library Comparison¶

In lieu of an external GNN library dependency, the unit-tested ml_first_principles.gnn_models package is the pinned reference: we copy our weights into its layers and check that the outputs match.

In [7]:

from ml_first_principles.gnn_models import GATLayer as GATRef, GCNLayer as GCNRef

# GCN: copy our weight matrix into the package layer and compare outputs.
# Our layer consumes the pre-normalized A_norm; the reference normalizes raw A itself
# (adding self-loops where the diagonal is zero), so both see the same operator.
gcn = GCNLayer(N, 16, rng)
gcn_ref = GCNRef(N, 16, random_state=SEED)
gcn_ref.weight = gcn.W
assert np.allclose(gcn.forward(X, A_norm), gcn_ref.forward(X, A), atol=1e-8)

# GAT: our single-head architecture (weight matrix W + attention vector a) matches
# the reference exactly, so we copy both parameters and compare.
gat = GATLayer(N, 16, rng)
gat_ref = GATRef(N, 16, random_state=SEED)
gat_ref.weight = gat.W
gat_ref.a = gat.a
out_gat, _ = gat.forward(X, A)
assert np.allclose(out_gat, gat_ref.forward(X, A), atol=1e-8)

print("GCN and GAT outputs match the unit-tested package reference.")

from ml_first_principles.gnn_models import GATLayer as GATRef, GCNLayer as GCNRef # GCN: copy our weight matrix into the package layer and compare outputs. # Our layer consumes the pre-normalized A_norm; the reference normalizes raw A itself # (adding self-loops where the diagonal is zero), so both see the same operator. gcn = GCNLayer(N, 16, rng) gcn_ref = GCNRef(N, 16, random_state=SEED) gcn_ref.weight = gcn.W assert np.allclose(gcn.forward(X, A_norm), gcn_ref.forward(X, A), atol=1e-8) # GAT: our single-head architecture (weight matrix W + attention vector a) matches # the reference exactly, so we copy both parameters and compare. gat = GATLayer(N, 16, rng) gat_ref = GATRef(N, 16, random_state=SEED) gat_ref.weight = gat.W gat_ref.a = gat.a out_gat, _ = gat.forward(X, A) assert np.allclose(out_gat, gat_ref.forward(X, A), atol=1e-8) print("GCN and GAT outputs match the unit-tested package reference.")

GCN and GAT outputs match the unit-tested package reference.

6. Experiments and Failures — VERIFY¶

In [8]:

# 2-Layer GCN Forward Pass (untrained, visualization of initialized features)
gcn1 = GCNLayer(N, 16, rng)
gcn2 = GCNLayer(16, 2, rng)  # 2 classes for visualization

H1 = relu(gcn1.forward(X, A_norm))
H2 = gcn2.forward(H1, A_norm)

plt.figure(figsize=(8, 6))
plt.scatter(H2[:, 0], H2[:, 1], c=labels, cmap='coolwarm', s=100, edgecolors='k')
plt.title("Untrained 2-Layer GCN Node Embeddings")
plt.xlabel("Dim 1")
plt.ylabel("Dim 2")
plt.show()

# 2-Layer GCN Forward Pass (untrained, visualization of initialized features) gcn1 = GCNLayer(N, 16, rng) gcn2 = GCNLayer(16, 2, rng) # 2 classes for visualization H1 = relu(gcn1.forward(X, A_norm)) H2 = gcn2.forward(H1, A_norm) plt.figure(figsize=(8, 6)) plt.scatter(H2[:, 0], H2[:, 1], c=labels, cmap='coolwarm', s=100, edgecolors='k') plt.title("Untrained 2-Layer GCN Node Embeddings") plt.xlabel("Dim 1") plt.ylabel("Dim 2") plt.show()

Training the GCN (Semi-Supervised Node Classification)¶

Now we actually train the 2-layer GCN with manual analytic gradients. Only two nodes are labeled — the instructor (node 0) and the president (node 33) — and softmax cross-entropy is applied to those rows only. With $S = \tilde{D}^{-1/2}\tilde{A}\tilde{D}^{-1/2}$ the forward pass is

$$H_1 = \mathrm{ReLU}(S X W_1), \qquad Z = S H_1 W_2,$$

so backpropagation is a plain matmul chain:

$$\frac{\partial L}{\partial W_2} = (S H_1)^\top \frac{\partial L}{\partial Z}, \qquad \frac{\partial L}{\partial H_1} = S^\top \frac{\partial L}{\partial Z} W_2^\top, \qquad \frac{\partial L}{\partial W_1} = (S X)^\top \left(\frac{\partial L}{\partial H_1} \odot \mathbf{1}[S X W_1 > 0]\right).$$

In [9]:

# Semi-supervised training: only nodes 0 (Mr. Hi) and 33 (Officer) are labeled
S = A_norm
labeled = np.array([0, 33])
y_onehot = np.zeros((labeled.size, 2))
y_onehot[np.arange(labeled.size), labels[labeled]] = 1.0

hidden_dim = 8
W1 = rng.normal(0, 0.1, (N, hidden_dim))
W2 = rng.normal(0, 0.1, (hidden_dim, 2))

lr = 0.5
epochs = 300
train_losses = []
SX = S @ X

for epoch in range(epochs):
    # Forward
    M = SX @ W1
    H1 = relu(M)
    SH1 = S @ H1
    Z = SH1 @ W2
    P = softmax(Z, axis=1)
    train_losses.append(-np.mean(np.log(P[labeled, labels[labeled]] + 1e-12)))

    # Backward: softmax cross-entropy on the labeled rows only
    dZ = np.zeros_like(Z)
    dZ[labeled] = (P[labeled] - y_onehot) / labeled.size
    dW2 = SH1.T @ dZ
    dH1 = S.T @ dZ @ W2.T
    dW1 = SX.T @ (dH1 * (M > 0))

    W1 -= lr * dW1
    W2 -= lr * dW2

plt.figure(figsize=(8, 5))
plt.plot(train_losses)
plt.title("GCN Semi-Supervised Training (2 labeled nodes)")
plt.xlabel("Epoch")
plt.ylabel("Cross-entropy on labeled nodes")
plt.grid(True)
plt.show()

# Evaluate against the true club split
P = softmax(S @ relu(SX @ W1) @ W2, axis=1)
pred = np.argmax(P, axis=1)
labeled_acc = np.mean(pred[labeled] == labels[labeled])
overall_acc = np.mean(pred == labels)
print(f"Labeled-node accuracy: {labeled_acc:.2f}, overall accuracy: {overall_acc:.4f}")
assert labeled_acc == 1.0
assert overall_acc >= 0.9, "Trained GCN should recover the two-community split"

# Semi-supervised training: only nodes 0 (Mr. Hi) and 33 (Officer) are labeled S = A_norm labeled = np.array([0, 33]) y_onehot = np.zeros((labeled.size, 2)) y_onehot[np.arange(labeled.size), labels[labeled]] = 1.0 hidden_dim = 8 W1 = rng.normal(0, 0.1, (N, hidden_dim)) W2 = rng.normal(0, 0.1, (hidden_dim, 2)) lr = 0.5 epochs = 300 train_losses = [] SX = S @ X for epoch in range(epochs): # Forward M = SX @ W1 H1 = relu(M) SH1 = S @ H1 Z = SH1 @ W2 P = softmax(Z, axis=1) train_losses.append(-np.mean(np.log(P[labeled, labels[labeled]] + 1e-12))) # Backward: softmax cross-entropy on the labeled rows only dZ = np.zeros_like(Z) dZ[labeled] = (P[labeled] - y_onehot) / labeled.size dW2 = SH1.T @ dZ dH1 = S.T @ dZ @ W2.T dW1 = SX.T @ (dH1 * (M > 0)) W1 -= lr * dW1 W2 -= lr * dW2 plt.figure(figsize=(8, 5)) plt.plot(train_losses) plt.title("GCN Semi-Supervised Training (2 labeled nodes)") plt.xlabel("Epoch") plt.ylabel("Cross-entropy on labeled nodes") plt.grid(True) plt.show() # Evaluate against the true club split P = softmax(S @ relu(SX @ W1) @ W2, axis=1) pred = np.argmax(P, axis=1) labeled_acc = np.mean(pred[labeled] == labels[labeled]) overall_acc = np.mean(pred == labels) print(f"Labeled-node accuracy: {labeled_acc:.2f}, overall accuracy: {overall_acc:.4f}") assert labeled_acc == 1.0 assert overall_acc >= 0.9, "Trained GCN should recover the two-community split"

Labeled-node accuracy: 1.00, overall accuracy: 0.9706

In [10]:

# Failure Case: Over-smoothing
# To isolate propagation from random projections, repeatedly apply the SAME
# normalized operator S = D^{-1/2} (A+I) D^{-1/2} with identity weights:
# any convergence is then due to propagation alone.

num_layers = 15
h = X.copy()
similarities = []

for _ in range(num_layers):
    h = A_norm @ h  # one propagation step, no learned weights
    
    # Compute average pairwise cosine similarity
    norm_h = h / (np.linalg.norm(h, axis=1, keepdims=True) + 1e-8)
    sim_matrix = norm_h @ norm_h.T
    
    # Average similarity of off-diagonal elements
    mask = ~np.eye(N, dtype=bool)
    similarities.append(np.mean(sim_matrix[mask]))

plt.figure(figsize=(8, 5))
plt.plot(range(1, num_layers + 1), similarities, marker='o')
plt.title(r"Over-smoothing: Repeated Propagation with $\hat{A}$")
plt.xlabel("Number of Propagation Steps")
plt.ylabel("Average Pairwise Cosine Similarity")
plt.grid(True)
plt.show()

# Deterministic check (no randomness involved): depth ~10 is far more smoothed than depth 1
print(f"Similarity at depth 1: {similarities[0]:.3f}, at depth 10: {similarities[9]:.3f}")
assert similarities[9] > similarities[0] + 0.5, "Expected clear over-smoothing by depth 10"

# Failure Case: Over-smoothing # To isolate propagation from random projections, repeatedly apply the SAME # normalized operator S = D^{-1/2} (A+I) D^{-1/2} with identity weights: # any convergence is then due to propagation alone. num_layers = 15 h = X.copy() similarities = [] for _ in range(num_layers): h = A_norm @ h # one propagation step, no learned weights # Compute average pairwise cosine similarity norm_h = h / (np.linalg.norm(h, axis=1, keepdims=True) + 1e-8) sim_matrix = norm_h @ norm_h.T # Average similarity of off-diagonal elements mask = ~np.eye(N, dtype=bool) similarities.append(np.mean(sim_matrix[mask])) plt.figure(figsize=(8, 5)) plt.plot(range(1, num_layers + 1), similarities, marker='o') plt.title(r"Over-smoothing: Repeated Propagation with $\hat{A}$") plt.xlabel("Number of Propagation Steps") plt.ylabel("Average Pairwise Cosine Similarity") plt.grid(True) plt.show() # Deterministic check (no randomness involved): depth ~10 is far more smoothed than depth 1 print(f"Similarity at depth 1: {similarities[0]:.3f}, at depth 10: {similarities[9]:.3f}") assert similarities[9] > similarities[0] + 0.5, "Expected clear over-smoothing by depth 10"

Similarity at depth 1: 0.138, at depth 10: 0.883

In [11]:

# Visualize GAT attention weights
gat = GATLayer(N, 16, rng)
out, att = gat.forward(X, A)

plt.figure(figsize=(8, 6))
plt.imshow(att, cmap='viridis')
plt.colorbar()
plt.title("GAT Attention Weights Heatmap")
plt.xlabel("Target Node")
plt.ylabel("Source Node")
plt.show()

# Visualize GAT attention weights gat = GATLayer(N, 16, rng) out, att = gat.forward(X, A) plt.figure(figsize=(8, 6)) plt.imshow(att, cmap='viridis') plt.colorbar() plt.title("GAT Attention Weights Heatmap") plt.xlabel("Target Node") plt.ylabel("Source Node") plt.show()

7. Connections¶

• The GCN normalization $\hat{A}$ acts as a low-pass filter on the graph signal, which is why stacking many layers causes over-smoothing.
• GAT's attention mechanism is mathematically similar to self-attention in Transformers, but restricted to the local neighborhood defined by the adjacency matrix.