Message Passing Framework | Graph Machine Learning

1. Why Message Passing?

We cannot feed a graph into an MLP as a raw adjacency matrix because:

Graphs have different sizes ($n$ varies)
There is no fixed node order - node 3 in one labeling might be node 7 in another
What matters is structure, not the arbitrary IDs

Message passing solves this by processing each node locally using its neighbors - the same rule applies regardless of graph size or node naming.

Intuition: Each node asks its neighbors "what do you know?" and combines their answers to update its own belief. After $L$ rounds, information has spread $L$ hops.

Refer: Gilmer et al. - Neural Message Passing for Quantum Chemistry (MPNN)

2. Permutation Invariance and Equivariance

Let $\mathbf{P}$ be any permutation matrix (reorders nodes). It reorders adjacency: $\mathbf{A}' = \mathbf{P}\mathbf{A}\mathbf{P}^T$ and features: $\mathbf{X}' = \mathbf{P}\mathbf{X}$.

Permutation equivariance (node-level output)

A function $f$ is equivariant if reordering input reorders output the same way:

$$f(\mathbf{P}\mathbf{A}\mathbf{P}^T, \mathbf{P}\mathbf{X}) = \mathbf{P}\, f(\mathbf{A}, \mathbf{X})$$

Each GNN layer should be equivariant - node $i$'s embedding depends on its neighborhood structure, not its index.

Permutation invariance (graph-level output)

A function $g$ is invariant if:

$$g(\mathbf{P}\mathbf{A}\mathbf{P}^T, \mathbf{P}\mathbf{X}) = g(\mathbf{A}, \mathbf{X})$$

Graph classification needs invariance - achieved by readout (sum/mean over nodes - symmetric function).

Why sum aggregation is invariant: $\sum_i \mathbf{h}_i = \sum_i (\mathbf{P}\mathbf{h})_i$ because permutation only reorders terms in the sum. Same for mean and max.

3. The MPNN Framework

Gilmer et al. (2017) unified GNNs into three steps at layer $l$:

Step 1 - Message

$$\mathbf{m}_{ij}^{(l+1)} = \phi^{(l)}\left(\mathbf{h}_i^{(l)}, \mathbf{h}_j^{(l)}, \mathbf{e}_{ij}\right)$$

Build a message from sender $j$ to receiver $i$ using both embeddings and optional edge features.

Step 2 - Aggregate

$$\bar{\mathbf{m}}_i^{(l+1)} = \bigoplus_{j \in \mathcal{N}(i)} \mathbf{m}_{ij}^{(l+1)}$$

$\bigoplus$ is a permutation-invariant aggregator: sum, mean, max, or attention-weighted sum.

Step 3 - Update

$$\mathbf{h}_i^{(l+1)} = \gamma^{(l)}\left(\mathbf{h}_i^{(l)}, \bar{\mathbf{m}}_i^{(l+1)}\right)$$

Combine old embedding with aggregated messages. Often an MLP or GRU.

Initial embedding $\mathbf{h}_i^{(0)} = \mathbf{x}_i$ (raw features).

Model	Message $\phi$	Aggregate $\bigoplus$	Update $\gamma$
GCN	normalized $h_j$	sum	linear + ReLU
GraphSAGE	$h_j$	mean/max/LSTM	concat + MLP
GAT	attention-weighted $h_j$	weighted sum	linear + ELU
GIN	$h_j$	sum	MLP$(h_i + \sum h_j)$

Full details per model: Classic Architectures.

4. Receptive Field and Depth

After $L$ message passing layers, node $i$'s embedding depends on its $L$-hop neighborhood (all nodes within $L$ edges).

Intuition: Layer 1 = direct neighbors. Layer 2 = neighbors of neighbors. Deeper = wider context but also more oversmoothing risk.

Why not 100 layers? On most graphs, 2–4 layers work best. Deeper layers mix all nodes toward similar embeddings (oversmoothing) and suffer over-squashing (bottleneck edges lose information). See Deep GNNs & Limitations.

5. Graph-Level Readout

For graph classification, pool node embeddings into one vector:

$$\mathbf{h}_G = \text{READOUT}\left(\{\mathbf{h}_i^{(L)} : i \in V\}\right)$$

Common READOUT functions:

Sum: $\sum_i \mathbf{h}_i$ - preserves size information (GIN uses this)
Mean: $\frac{1}{n}\sum_i \mathbf{h}_i$ - size-invariant
Max: $\max_i \mathbf{h}_i$ per dimension - captures strongest signal
Attention pooling: learn weights per node
Jumping Knowledge: combine embeddings from multiple layers

6. Matrix Form View

One GCN-style layer can be written for all nodes at once:

$$\mathbf{H}^{(l+1)} = \sigma\left(\tilde{\mathbf{D}}^{-1/2}\tilde{\mathbf{A}}\tilde{\mathbf{D}}^{-1/2}\mathbf{H}^{(l)}\mathbf{W}^{(l)}\right)$$

where $\tilde{\mathbf{A}} = \mathbf{A} + \mathbf{I}$. This is exactly "multiply by adjacency (spread messages), then linear transform." Sparse matrix multiply makes this efficient.

7. Batching Multiple Graphs

PyG batches graphs by creating one disjoint union - stack all nodes and edges, offset edge indices, add a batch vector labeling which graph each node belongs to.

from torch_geometric.loader import DataLoader
loader = DataLoader(dataset, batch_size=32, shuffle=True)
for batch in loader:
    out = model(batch.x, batch.edge_index, batch.batch)

For single large graphs (Cora, OGBN-Arxiv), use NeighborLoader to sample local subgraphs instead.

8. PyG Implementation

from torch_geometric.nn import MessagePassing

class CustomMP(MessagePassing):
    def __init__(self, in_ch, out_ch):
        super().__init__(aggr='add')  # sum aggregation
        self.lin = torch.nn.Linear(in_ch, out_ch)

    def forward(self, x, edge_index):
        return self.propagate(edge_index, x=x)

    def message(self, x_j):
        # x_j: features of source nodes j for each edge
        return x_j

    def update(self, aggr_out):
        return self.lin(aggr_out).relu()

message defines $\phi$, aggr defines $\bigoplus$, update defines $\gamma$.