Implement Graph Convolution Network (GCN) Layer

Problem

Implement a single Graph Convolutional Network (GCN) layer from scratch. Given a node feature matrix X, an adjacency matrix A, and a weight matrix W, apply the GCN propagation rule: add self-loops, compute the symmetric normalization of the adjacency matrix, and multiply through.

Solution

def gcn_layer(X: list, A: list, W: list) -> list:
    n = len(A)
    A_hat = [[A[i][j] + (1 if i == j else 0) for j in range(n)] for i in range(n)]
    d_inv_sqrt = [1.0 / sum(A_hat[i])**0.5 for i in range(n)]
    A_norm = [[d_inv_sqrt[i] * A_hat[i][j] * d_inv_sqrt[j] for j in range(n)] for i in range(n)]
    f, fp = len(X[0]), len(W[0])
    AX = [[sum(A_norm[i][k] * X[k][j] for k in range(n)) for j in range(f)] for i in range(n)]
    return [[sum(AX[i][k] * W[k][j] for k in range(f)) for j in range(fp)] for i in range(n)]

Explanation

1. Self-loops: Add the identity matrix to A so each node aggregates its own features along with its neighbors'.
2. Degree matrix: Compute D_hat where each diagonal entry is the sum of the corresponding row of A_hat.
3. Symmetric normalization: Compute D^{-1/2} A_hat D^{-1/2} which normalizes messages by the geometric mean of the sender's and receiver's degrees.
4. Propagation: Multiply the normalized adjacency by features X and then by the learnable weight matrix W.

Complexity

• Time: O(n^2 f + n f * f') where n is nodes, f is input features, f' is output features
• Space: O(n^2) for the adjacency and normalization matrices