Graph Neural Networks

Graph Neural Networks (GNNs) learn node, edge, and graph level representations by passing messages along the edges of a graph. Each layer lets a node gather information from its neighbors, so stacking layers grows the receptive field one hop at a time.

Message Passing Framework

Almost every GNN is an instance of the message passing scheme. For node $v$ with neighbors $\mathcal{N}(v)$ and layer $k$:

$$\begin{array}{rl}{m}_v^{(k)}& =\text{AGG}(\{\varphi (h_u^{(k-1)},h_v^{(k-1)},e_{uv}):u\in \mathcal{N}(v)\})\\ h_v^{(k)}& =\psi (h_v^{(k-1)},m_v^{(k)})\end{array}$$

• $\varphi$ builds a message from a source node, target node, and edge feature.
• $\text{AGG}$ is a permutation invariant aggregator (sum, mean, max, or attention) so the output does not depend on neighbor ordering.
• $\psi$ updates the node using its old state and the aggregated message.

Common Architectures

Model	Aggregation	Key idea
GCN	normalized mean	symmetric degree normalization ${\tilde{D}}^{-1/2}\tilde{A}{\tilde{D}}^{-1/2}$
GraphSAGE	mean / max / LSTM	sample a fixed number of neighbors for scalability
GAT	attention weighted sum	learn per edge weights with Self-Attention style scores

GCN uses the whole (normalized) adjacency, GraphSAGE samples neighbors so it scales to huge graphs, and GAT computes attention coefficients ${\alpha }_{uv}$ so important neighbors count more.

PyTorch: A GCN Layer

Here we implement one GCN propagation using the dense normalized adjacency. torch.einsum makes the neighbor aggregation explicit as a matrix product over the node axis.

import torch
import torch.nn as nn
 
class GCNLayer(nn.Module):
    def __init__(self, d_in, d_out):
        super().__init__()
        self.lin = nn.Linear(d_in, d_out, bias=False)
 
    def forward(self, x, a_hat):
        # x:     (n, d_in)  node features
        # a_hat: (n, n)     normalized adjacency D^-1/2 (A + I) D^-1/2
        h = self.lin(x)                       # (n, d_out) transform first
        # aggregate: each node = weighted sum of neighbor messages
        out = torch.einsum("nm,md->nd", a_hat, h)
        return torch.relu(out)
 
 
def normalize_adjacency(a):
    n = a.shape[0]
    a_tilde = a + torch.eye(n, device=a.device)     # add self loops
    deg = a_tilde.sum(dim=-1)                        # (n,)
    d_inv_sqrt = deg.pow(-0.5)
    # D^-1/2 (A+I) D^-1/2 via outer product of degree factors
    return torch.einsum("i,ij,j->ij", d_inv_sqrt, a_tilde, d_inv_sqrt)

The einsum("i,ij,j->ij", ...) scales row $i$ and column $j$ of $\tilde{A}$ by their inverse square root degrees in one shot.

Over-Smoothing

Over-smoothing

Stacking many message passing layers repeatedly averages neighborhoods, so node embeddings converge to nearly identical vectors and lose discriminative power. This is why most GNNs are shallow (2 to 4 layers). Residual connections, LayerNorm, jumping knowledge connections, and PairNorm mitigate it, but deep GNNs remain hard to train.

Practical Notes

• Sum aggregation is the most expressive (it can count neighbors), which is the basis of the Graph Isomorphism Network; mean and max lose multiplicity information.
• For large graphs use neighbor sampling (GraphSAGE) or subgraph batching; the dense adjacency above is only practical for small graphs.
• Edge features and global readout (sum or mean over all nodes) extend the same framework to edge and graph level prediction.