Alessio Borgi

Alessio Borgi

PhD in Graph Neural Networks and Generative AI 👨🏻‍💻🤖

Message Passing: The Universal GNN Framework

5 minute read

Published:

TL;DR: Message Passing Neural Networks (Gilmer et al., 2017) provide a unified framework for all GNNs. Each layer runs three steps: MESSAGE (what each neighbour sends), AGGREGATE (collect all messages), UPDATE (compute new node representation). Choosing different functions for each step gives you different GNN architectures.
Message Passing Neural Network
Message Passing Neural Network (MPNN) framework (Gilmer et al., 2017)

The Framework

The MPNN framework (Gilmer et al., 2017, ICML) defines GNN computation through a series of message passing steps. At each step \(t\):

\[ m_v^{(t+1)} \;=\; \operatorname{AGGREGATE}\Big(\big\{\, \operatorname{MSG}\big(h_v^{(t)},\, h_u^{(t)},\, e_{uv}\big) \;:\; u \in \mathcal{N}(v) \,\big\}\Big) \] \[ h_v^{(t+1)} \;=\; \operatorname{UPDATE}\big(h_v^{(t)},\, m_v^{(t+1)}\big) \]

Where:

  • \(h_v^{(t)} \in \mathbb{R}^{d_t}\) — the representation of node \(v\) after \(t\) message passing steps; \(h_v^{(0)}\) is the input feature vector.
  • \(\mathcal{N}(v)\) — the set of neighbours of \(v\) in the graph.
  • \(e_{uv}\) — the (optional) feature vector of the edge between \(u\) and \(v\).
  • \(m_v^{(t+1)}\) — the aggregated message arriving at \(v\) at step \(t+1\).
  • \(\operatorname{MSG}\) — the message function, computing what a neighbour sends.
  • \(\operatorname{AGGREGATE}\) — combines all incoming messages (must be permutation-invariant, since \(\{\cdot\}\) is a multiset, not an ordered list).
  • \(\operatorname{UPDATE}\) — computes the new representation from the old one plus the aggregated message.
B h_B A h_A MSG(h_A, h_B) C h_C MSG(h_C, h_B) D h_D MSG(h_D, h_B) ② AGGREGATE ③ UPDATE new h_B ✓ ① Compute messages
Figure 1: Node B receives messages from its three neighbours A, C, D. The messages are aggregated (e.g., summed or averaged), then combined with B's own representation in an UPDATE function to produce a new \(h_B\).

Concrete Worked Example: One Full MPNN Step

Let node \(B\) have features \(h_B = [1, 0]\), with three neighbours \(\mathcal{N}(B) = \{A, C, D\}\) whose features are \(h_A = [0, 1]\), \(h_C = [1, 1]\) and \(h_D = [0, 0]\).

Step 1 — Compute messages (using the identity message function, \(\operatorname{MSG}(h_v, h_u, e_{uv}) = h_u\), i.e. just pass the neighbour’s features along):

\[ m_{A \to B} = [0, 1], \qquad m_{C \to B} = [1, 1], \qquad m_{D \to B} = [0, 0] \]

Step 2 — Aggregate (sum):

\[ m_B = [0,1] + [1,1] + [0,0] = [1, 2] \]

Step 3 — Update (concatenate own features with the aggregate, apply a learned linear map \(W \in \mathbb{R}^{2 \times 4}\) and a ReLU):

\[ h_B' = \operatorname{ReLU}\big(W \, [\, h_B \,\Vert\, m_B \,]\big) = \operatorname{ReLU}\big(W \, [1, 0, 1, 2]^{\top}\big) \]

Here \(\Vert\) denotes concatenation, so \([\, h_B \Vert m_B \,] \in \mathbb{R}^4\). After this one layer, \(B\)’s new 2-dimensional embedding encodes information from all three of its neighbours.

Key Insight: The three steps — MSG, AGGREGATE, UPDATE — are independent design choices. Changing any one of them gives a different GNN family. GCN sends \(W h_u\) and aggregates with a degree-normalised sum. GAT weights that sum by learned attention coefficients. GIN uses a plain sum followed by an MLP. The framework shows that these are all variations on the same theme.

Step 1: Message Function

The message function computes what each neighbour sends. The simplest choice: just send the neighbour’s features.

\[ \begin{aligned} \operatorname{MSG}(h_v, h_u, e_{uv}) &= h_u && \text{(pass the raw neighbour features)}\\ \operatorname{MSG}(h_v, h_u, e_{uv}) &= W h_u && \text{(GCN: linear transform first)}\\ \operatorname{MSG}(h_v, h_u, e_{uv}) &= \alpha_{vu} \, W h_u && \text{(GAT: scale by an attention weight)} \end{aligned} \]

Here \(W\) is a learned weight matrix shared by all edges, and \(\alpha_{vu}\) is the scalar attention weight GAT places on the edge \(u \to v\). Including edge features \(e_{uv}\) allows the model to distinguish bond types in a molecule or relationship types in a knowledge graph.

Step 2: Aggregate Function

The aggregation combines all messages. It must be permutation-invariant (the order of neighbours shouldn’t matter):

AggregatorFormulaProperties
Sum\(\sum_{u \in \mathcal{N}(v)} m_{u \to v}\)Keeps the size of the neighbourhood
Mean\(\frac{1}{\lvert \mathcal{N}(v) \rvert} \sum_{u \in \mathcal{N}(v)} m_{u \to v}\)Normalised, size-invariant
Max\(\max_{u \in \mathcal{N}(v)} m_{u \to v}\) (elementwise)Captures the most extreme feature
Attention-weighted\(\sum_{u \in \mathcal{N}(v)} \alpha_{vu} \, m_{u \to v}\)Adaptive, like GAT

GIN (see the GIN post) shows that sum is the most expressive of these: it is an injective function of the multiset of messages when the message space is countable, so it distinguishes every neighbourhood that the 1-WL test distinguishes. Mean and max are not injective — mean discards the neighbourhood size, max discards multiplicities — so both lose structural information.

Step 3: Update Function

Given the aggregated message \(m_v\) and the old representation \(h_v\), compute the new one:

\[ \begin{aligned} h_v' &= \sigma\big(W \, [\, h_v \Vert m_v \,]\big) && \text{(single linear layer + non-linearity)}\\ h_v' &= \operatorname{GRU}\big(h_v,\, m_v\big) && \text{(recurrent update, as in the original MPNN)}\\ h_v' &= \operatorname{MLP}\big([\, h_v \Vert m_v \,]\big) && \text{(GraphSAGE-style)} \end{aligned} \]

where \(\sigma\) is an elementwise non-linearity (usually ReLU) and \([\,\cdot \Vert \cdot\,]\) is concatenation.

A Running Example: Molecule Property Prediction

Consider predicting if a molecule is toxic:

  • Nodes = atoms (features: atom type, charge, is_aromatic)
  • Edges = bonds (features: bond type: single/double/triple)
  • After \(k\) MPNN layers, each atom knows about its \(k\)-hop neighbourhood.
  • A readout \(h_G = R(\{h_v^{(k)} : v \in V\})\) aggregates all atom embeddings into a single graph embedding; \(R\) must itself be permutation-invariant.
  • An MLP predicts toxicity from \(h_G\).

After 3 layers, an atom “knows” about the atoms 3 bonds away — capturing local chemical environments like functional groups.

✅ Key Takeaways

  • All GNNs are instances of MPNN: choose the \(\operatorname{MSG}\), \(\operatorname{AGGREGATE}\), and \(\operatorname{UPDATE}\) functions.
  • \(\operatorname{AGGREGATE}\) must be permutation-invariant. Sum is the most expressive choice, because it is injective over multisets (GIN).
  • After \(k\) layers, each node's embedding captures its \(k\)-hop neighbourhood.
  • Graph-level predictions require a readout function that pools node embeddings into a single vector.

References

You May Also Enjoy

The Standard Library Modules Worth Knowing Before You pip install

8 minute read

Published:

Python ships with a counter, a queue, a memoiser, a CLI parser, a test framework’s smaller cousin, and a JSON-safe date library already installed — most of what people reach for a package to solve is one import away.

Files and Context Managers: Why with Is Not Optional

9 minute read

Published:

A file object that falls out of scope without with does get closed eventually — but ‘eventually’ means whenever the garbage collector gets around to it, which is not a promise any program handling more than a handful of files can live with.

Errors and Exceptions: EAFP, the Hierarchy, and Reading a Traceback

9 minute read

Published:

Exceptions in Python are not exceptional. They are a normal control-flow mechanism that the language leans on so heavily that the idiomatic style is to try the operation and handle the failure, rather than check first — which is faster, shorter, and free of race conditions.