Message Passing Neural Networks

• Message Passing Neural Networks (MPNNs) are graph neural architectures that iteratively pass and update node and edge features to produce rich representations.
• They employ differentiable message, update, and readout functions, ensuring permutation invariance and effective integration of structural graph information.
• MPNNs have demonstrated state-of-the-art performance in molecular machine learning and physical simulations, leveraging scalable design innovations.

Message Passing Neural Networks (MPNNs) are a family of neural architectures designed to operate on graph-structured data by iteratively exchanging, transforming, and aggregating features among nodes and edges to produce expressive representations at both local and global scales. Originally formulated to unify prior variants of graph neural networks, the MPNN framework is distinguished by its invariance to graph isomorphism and its ability to incorporate both node and edge attributes within a principled, flexible propagation and aggregation workflow (Gilmer et al., 2017). MPNNs have become a cornerstone in molecular machine learning, physical simulation, network science, and numerous domains where relational and structured information is fundamental.

1. Mathematical Framework and Formalism

The MPNN paradigm abstractly defines a computation over a graph $G = (V, E)$ with node features $x_v$ and edge features $e_{vw}$. The architecture is organized as follows (Gilmer et al., 2017):

• Message Passing Phase (for $T$ steps):

$$\begin{align} m_v^{t+1} &= \sum_{w \in N(v)} M_t(h_v^t, h_w^t, e_{vw}) \ h_v^{t+1} &= U_t(h_v^t, m_v^{t+1}) \end{align}$$

• $h_v^t$: hidden state (feature vector) of node $v$ at iteration $t$
• $M_t$: message function, differentiable with respect to its arguments; responsible for encoding the way features from neighbors and edge attributes are combined
• $U_t$: vertex update function, typically a small neural network (e.g., GRU, MLP), responsible for integrating current state and incoming message
• $N(v)$: neighborhood of $v$
  • Readout Phase:

$\hat{y} = R\left(\{h_v^T \mid v \in G\}\right)$

• $R$: readout function, must be permutation-invariant (e.g., sum, mean, set2set). Produces graph-level (or node/edge-level) output.

This abstraction encompasses numerous GNN designs—including early graph convolutional networks, gated graph neural networks, and neural fingerprints—by selecting appropriate forms for $M_t$, $U_t$, and $R$. The architecture is designed to ensure invariance under permutations of node ordering—a critical requirement for graph isomorphism invariance.

2. Architectural Variants and Extensions

Gilmer et al. introduced several MPNN variations targeting improved expressiveness, scalability, and data efficiency (Gilmer et al., 2017).

Edge Network Message Function: Enables incorporation of continuous or high-dimensional edge attributes (e.g., distances) through a neural network $A(e_{vw})$ that outputs a matrix to transform neighbor features before aggregation: $M(h_v, h_w, e_{vw}) = A(e_{vw}) h_w$

Pair Message Function: Jointly conditions on both sender and receiver node states and the edge: $M(h_v, h_w, e_{vw}) = f(h_w, h_v, e_{vw})$ where $f$ is a neural network.

Virtual Graph Elements:

• Virtual edges enable information flow between distant nodes by adding synthetic edges.
• Master node acts as a centralized aggregator, sharing information globally with all nodes.

Towers Variant: Improves scalability by partitioning the node state into $k$ blocks ("towers") and performing disjoint message passing per block, followed by mixing (via a shared neural network), reducing computational complexity from $O(n^2 d^2)$ to $O(n^2 d^2 / k)$.

Advanced Readout: The set2set model (an attention-based readout) aggregates node features into a richer graph representation.

3. Aggregation, Readout, and Permutation Invariance

Aggregation occurs in two places: (i) message collection—typically by sum (or mean), which is permutation-invariant; (ii) graph-level readout, which must also be invariant across node orderings. Simple summation is common, but attention or set-based mechanisms like set2set can be used to increase expressiveness. The invariance constraints are essential for models to retain physical and structural meaning, especially in chemical applications where node permutations do not alter the atomistic system (Gilmer et al., 2017).

4. Benchmarks and Empirical Performance

MPNNs exhibited state-of-the-art performance on the QM9 quantum property prediction benchmark, comprising ~134k molecules labeled with 13 properties. The best variant—using edge networks, set2set readout, explicit hydrogens, and ensembling—achieved error ratios below 1 (chemical accuracy) on 11 out of 13 properties and outperformed all prior baselines (Gilmer et al., 2017).

Target	Best Previous	MPNN (single)	MPNN Ensemble
Average Error Ratio	2.17 (BAML), 2.08 (BOB)	0.68	0.52

Ablation studies revealed that incorporating explicit spatial information, virtual elements, and set-based readouts led to measurably improved accuracy. MPNNs showed superior data efficiency, matching or outperforming baselines even with smaller training datasets.

The towers variant facilitated efficient scaling to higher-dimensional embeddings without quadratic computational scaling, further enabling use on larger molecules.

5. Scalability and Architectural Considerations

The scalability of MPNNs depends on message/aggregation strategies and hidden state dimension. The towers construction partitions node states, significantly limiting per-layer compute and memory costs. The ability to inject virtual edges or global nodes (e.g., master nodes) enables the architecture to efficiently capture long-range interactions otherwise inaccessible via standard local message passing. The computation can be parallelized by limiting the receptive field and careful batching of graphs.

Parameter-sharing and permutation-invariant aggregation allow entire batches of varying graphs to be processed simultaneously, essential for contemporary deep learning workflows.

6. Analysis, Limitations, and Future Directions

The formalism clarifies that the expressive power of MPNNs is upper-bounded by their ability to propagate information through node neighborhoods within $T$ steps, and is conditioned on both node and edge feature representations. Extensions such as multi-hop aggregation, virtual elements, or attention mechanisms offer paths to overcome locality bottlenecks. Careful handling of atomic hydrogen representation and explicit spatial features proves to be important in molecular applications.

Identified limitations include difficulty scaling to very large graphs due to $O(n^2)$ scaling in dense or long-range models, and potential constraints in distinguishing certain classes of nonisomorphic graphs (cf. Weisfeiler-Lehman expressiveness barriers in later works). Future research is suggested to focus on learning attention over messages, developing models for larger graphs, and improving generalization across different graph sizes (Gilmer et al., 2017).

7. Summary Table: MPNN Formalism

Component	Description
Nodes	$v$, with feature $x_v$
Edges	$(v, w)$, with feature $e_{vw}$
Hidden State	$h_v^t$ at time (layer) $t$
Message aggregation	$$m_v^{t+1} = \sum_{w \in N(v)} M_t(h_v^t, h_w^t, e_{vw})$$
State update	$h_v^{t+1} = U_t(h_v^t, m_v^{t+1})$
Readout	$\hat{y} = R(\{h_v^T\})$, permutation invariant
Message, Update, Readout	Differentiable modules (commonly neural networks)

---

MPNNs as defined by Gilmer et al. established a potent and flexible paradigm for graph-structured learning. Their systematic formalism, rich ablation studies, and architectural innovations constitute a foundation for ongoing advances in graph machine learning, including later developments in equivariant models, expressive power analyses, and scalable system implementations.