import marimo as mo
Neural Message Passing in pytorch-geometric¶
August 4 2026 Training Workshop: Introduction to Deep Graph Learning Ingo Scholtes, CAIDAS, Julius-Maximilians-Universität Würzburg (JMU), Germany
We introduce neural message passing, a fundamental building block of Graph Neural Networks, and show how we can efficiently implement it based on the geometric deep learning package pytorch-geometric, which is commonly referred to as pyg.
import seaborn as sns
import torch
import torch_geometric
from matplotlib import pyplot as plt
import pathpyG as pp
plt.style.use('default')
sns.set_style("whitegrid")
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print('Running on', device)
Running on cpu
We build our example networks with pathpyG. As we have seen in Basic pathpyG Concepts, each pathpyG graph stores its edge index as well as arbitrary node-, edge- and graph-level attributes as tensors in a torch_geometric.data.Data object, which we can directly use for message passing in pyG.
For the following example network with five nodes, we add two attributes x and y to the nodes. Following the usual convention, the attribute x contains a node feature tensor. In our simple toy example from the lecutre, we use a one-dimensional tensor that contains a single number. The y attribute contains a tensor with a node-level target variable. For the purpose of illustration, we use two target values per node.
n = pp.Graph.from_edge_list([('a', 'b'), ('b', 'c'), ('a', 'c'), ('c', 'd'), ('d', 'e'), ('e', 'c')]).to_undirected()
n.data.x = torch.tensor([[1.], [2.], [3.], [4.], [5.]])
n.data.y = torch.tensor([[0, 1], [1, 0], [1, 0], [0, 1], [0, 1]])
pp.plot(n, edge_color='gray');
Printing the Data object gives the shape of the node features, the edge index, and the target labels, which can be accessed via the respective attributes of the Data instance:
print(n.data)
print(n.data.x)
print(n.data.edge_index)
print(n.data.y)
Data(edge_index=[2, 12], num_nodes=5, x=[5, 1], y=[5, 2])
tensor([[1.],
[2.],
[3.],
[4.],
[5.]])
EdgeIndex([[0, 0, 1, 1, 2, 2, 2, 2, 3, 3, 4, 4],
[1, 2, 0, 2, 0, 1, 3, 4, 2, 4, 2, 3]], sparse_size=(5, 5), nnz=12,
sort_order=row, is_undirected=True)
tensor([[0, 1],
[1, 0],
[1, 0],
[0, 1],
[0, 1]])
Implementing Message Passing in pyG¶
We are now ready to generalize the idea of convolutional neural networks to graphs. In convolutional neural networks for image data, image convolutions are implemented by adding each pixel of an image to the pixel values of the neighboring pixels and additionally applying an appropriate weighting scheme. This can be generalized to networks, where we can apply a similar scheme to nodes and their neighbors.
This can be expressed as a message passing algorithm, where in each step of the convolution nodes exchange their current feature values with their neighbors. Each node then takes the incoming "messages" from its neighbors, aggregates them and computes a new local value. Depending on how much information from the neighborhood we wish to incorporate in the learning process, this convolution can be repeated multiple times (just as in convolutional neural networks).
While we could implement message passing on our own, torch_geometric conveniently provides a base class that simplifies this process (and makes it fast by calculating it at the GPU). The process of implementing a graph convolutional layer based on the MessagePassing base class is explained in detail here. Note that there are many different ways to implement the message passing (using different aggregation rules, weighting schemes, or normalizations) and depending on the specific implementation we obtain different graph neural networks.
Following the lecture, here we focus on the approach proposed by Kipf and Welling in 2017, which is a more or less direct translation of image convolutions to graphs. Using the notation from the lecture, the message passing process can be mathematically described via the following update rule for the node features $h_i^{(k)}$ in step $k$ of the message passing:
$$ h_i^{(k)} := \sigma \left(\mathbf{W}^{(k)}\sum_{j \in N_i \cup \{i\}} \frac{\left(h_j^{(k-1)}\right)}{\sqrt{d_i\cdot d_j}}\right)$$
Here $h_i^{(k)}$ is the feature of node $i$ in step $k$, $d_i$ is the degree of node $i$, $N_i$ is the set of neighbors of node $i$, $\mathbf{A}^{(k)}$ are learnable weight parameters and $\sigma$ is a non-linear activation function.
This can be viewed as a localized, graph-based generalization of the convolution commonly applied to image data (which can be viewed as a grid graph).
We will implement this step by step, starting with a simple aggregation of the neighbor features, i.e. we initially do not apply the symmetric degree-normalization as well as the multiplication with the weight matrix. We further do not include self-loops. We thus start with an update rule that simply adds the features of all neighbors.
$$ h_i^{(k)} :=\sum_{j \in N_i} h_j^{(k-1)}$$
In pyG we can implement this via a message passing class that derived from the torch_geometric.nn.MessagePassing base class. In the constructor of the base class, we can set the aggregation function that should be used by the nodes receiving messages (i.e. features) from their neighbors.
We basically have to implement two functions:
forwardis analogous to the forward function in atorch.nn.Module, i.e. here we specify the function that is applied when a tensorxwith associated edge indexedge_indexis passed as input. We call theself.propagatefunction, which will trigger a single round of message passing along all edges in the graph.messageconstructs the message to be sent along all directed edges $(j,i)$ from node $j$ to node $i$. For a tensor x that is passed as an argument to the propagate function, we can use special variables "x_i" and "x_j" to extract the tensor associated with nodes i and j for all edges (j,i). Hence, x_j and x_i have shape [m, d] where m is the number of edges and d is the dimensionality of the node features x
The simple update rule outlined above can be implemented as follows:
class MP(torch_geometric.nn.MessagePassing):
"""A message passing layer that sums the features of neighboring nodes."""
# Arguments passed to `propagate()`. pyG usually infers them from the source code of the class,
# which is not possible for classes defined in marimo notebooks
propagate_type = {'x': torch.Tensor}
def __init__(self):
"""Initialize the layer parameters."""
super().__init__(aggr='add')
def forward(self, x, edge_index):
"""Perform a single round of message passing."""
# this triggers a single round of message passing, where the message is passed along all edges in the network
return self.propagate(edge_index, x=x)
def message(self, x_j):
"""Construct the message that is sent along each edge."""
# This method constructs the message to be sent along all directed edges (j,i) from j to i
# For a tensor x that is passed as an argument to the propagate function, we can
# use special variables "x_i" and "x_j" to extract the tensor associated with nodes i
# and j for all edges (j,i). Hence, x_j and x_i have shape [m, d] where m is the number
# of edges and d is the dimensionality of the node features x
# Here we simply return a message that contains the feature of the source node j
print('Message = ', x_j.t())
return x_j
Let us now try this. Here is the feature vector that we use to initialize the node states in the message passing:
print(n.data.x)
tensor([[1.],
[2.],
[3.],
[4.],
[5.]])
Here is the graph:
pp.plot(n, edge_color='gray');
We now apply one round of message passing and print the result:
mp = MP()
_x = mp.forward(n.data.x, n.data.edge_index)
print(_x)
Message = tensor([[1., 1., 2., 2., 3., 3., 3., 3., 4., 4., 5., 5.]])
tensor([[ 5.],
[ 4.],
[12.],
[ 8.],
[ 7.]])
We can also use a mean aggregation, which corresponds to the following update rule:
$$ h_j^{(k)} = \frac{1}{d_i} \sum_{j \in N_i} h_j^{(k-1)}$$
class MP_1(torch_geometric.nn.MessagePassing):
"""A message passing layer that averages the features of neighboring nodes."""
# Arguments passed to `propagate()`. pyG usually infers them from the source code of the class,
# which is not possible for classes defined in marimo notebooks
propagate_type = {'x': torch.Tensor}
def __init__(self):
"""Initialize the layer parameters."""
super().__init__(aggr='mean')
def forward(self, x, edge_index):
"""Perform a single round of message passing."""
return self.propagate(edge_index, x=x) # this triggers a single round of message passing along all edges in the network
def message(self, x_j):
"""Construct the message that is sent along each edge."""
print('Message = ', x_j.t())
return x_j # This method constructs the message to be sent to nodes i for all edges (j,i) # For a tensor x that is passed as an argument to the propagate function, we can # use special variables "x_i" and "x_j" to extract the tensor associated with nodes i # and j for all edges (j,i). Hence, x_j and x_i have shape [m, d] where m is the number # of edges and d is the dimensionality of the node features x
mp_1 = MP_1()
_x = mp_1.forward(n.data.x, n.data.edge_index)
print(_x)
Message = tensor([[1., 1., 2., 2., 3., 3., 3., 3., 4., 4., 5., 5.]])
tensor([[2.5000],
[2.0000],
[3.0000],
[4.0000],
[3.5000]])
In both examples, we note that the nodes overwrite their own features, i.e. the values are replaced with the (aggregate) value received from the neighbors. To avoid this, we can add self-loops, which means that each node also sends a message to itself. This can be done by adding self-loops to the edge_index before running the message passing. We can use the convenience function torch_geometric.utils.add_self_loops as follows:
class MP_2(torch_geometric.nn.MessagePassing):
"""A message passing layer that sums the features of neighbors and of the node itself."""
# Arguments passed to `propagate()`. pyG usually infers them from the source code of the class,
# which is not possible for classes defined in marimo notebooks
propagate_type = {'x': torch.Tensor}
def __init__(self):
"""Initialize the layer parameters."""
super().__init__(aggr='add')
def forward(self, x, edge_index):
"""Perform a single round of message passing."""
edge_index, _ = torch_geometric.utils.add_self_loops(edge_index, num_nodes=x.size(0)) # add self-loops so that nodes include their own features in the aggregation
return self.propagate(edge_index, x=x)
def message(self, x_j): # this triggers a single round of message passing along all edges in the network
"""Construct the message that is sent along each edge."""
print('Message = ', x_j.t())
return x_j # This method constructs the message to be sent to nodes i for all edges (j,i) # For a tensor x that is passed as an argument to the propagate function, we can # use special variables "x_i" and "x_j" to extract the tensor associated with nodes i # and j for all edges (j,i). Hence, x_j and x_i have shape [m, d] where m is the number # of edges and d is the dimensionality of the node features x
We now find that the features of neighbours are added to the local features of the nodes:
mp_2 = MP_2()
_x = mp_2.forward(n.data.x, n.data.edge_index)
print(_x)
Message = tensor([[1., 1., 2., 2., 3., 3., 3., 3., 4., 4., 5., 5., 1., 2., 3., 4., 5.]])
tensor([[ 6.],
[ 6.],
[15.],
[12.],
[12.]])
We finally add the two missing ingredients for a graph convolutional layer, which is the normalization by the product of the square root of the node degrees as well as the linear transformation with learnable weights. For this, we first confirm our intuition that - so far - our message passing model actually has no parameters that we could fit to the data.
print([x for x in mp_2.parameters()])
[]
To understand the implementation of the convolution class, it is helpful to consider the degree-based normalization step by step. Let's first split the edge_index into two tensors that take the indices of source and target nodes:
source, target = n.data.edge_index
print(source)
print(target)
Index([0, 0, 1, 1, 2, 2, 2, 2, 3, 3, 4, 4], dim_size=5, is_sorted=True) Index([1, 2, 0, 2, 0, 1, 3, 4, 2, 4, 2, 3], dim_size=5)
We can use the torch_geometric.utils.degree to calculate node degrees based on a node index tensor. This method assumes that we pass a tensor that contains the source or target nodes of an edge index, which means that the in- or out-degree can simply be calculated by counting how often the indices occur.
For an undirected network, this method will calculate the degree of all nodes, independent of whether we pass the source or target nodes (since all links exist in both directions). So the following code calculates the degree sequence of our network:
deg = torch_geometric.utils.degree(target, n.data.x.size(0), dtype=torch.float32)
print(deg)
tensor([2., 2., 4., 2., 2.])
By indexing the degree tensor with the indices of source and target nodes, we can calculate the (in-)degrees of source and target nodes for all directed edges in the network.
print(deg[source])
print(deg[target])
tensor([2., 2., 2., 2., 4., 4., 4., 4., 2., 2., 2., 2.]) tensor([2., 4., 2., 4., 2., 2., 2., 2., 4., 2., 4., 2.])
Apart from an edge index and a tensor with the node features, the propagate function of the MessagePassing base class accepts an argument norm. It is a tensor with $m$ entries, one per edge, which capture the normalization factors that are applied in the exchange of all messages across all $m$ edges of the network. The symmetric degree-based normalization term
$$ \frac{1}{\sqrt{d_i} \cdot \sqrt{d_j}} $$
can be implemented as follows (additionally accounting for nodes with zero degree): The normalization factor for a directed edge (a,b) is
$$ \frac{1}{\sqrt{d_a} \cdot \sqrt{d_b}} = \frac{1}{\sqrt{2 \cdot 2}} := \frac{1}{2} $$
The normalization factor for directed edges (b,c) and (a,c) are
$$ \frac{1}{\sqrt{d_b} \cdot \sqrt{d_c}} = \frac{1}{\sqrt{4 \cdot 2}} = \frac{1}{2 \cdot \sqrt{2}} \approx 0.3536 $$
deg_inv_sqrt = deg.pow(-0.5)
deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0
norm = deg_inv_sqrt[source] * deg_inv_sqrt[target]
print(norm)
tensor([0.5000, 0.3536, 0.5000, 0.3536, 0.3536, 0.3536, 0.3536, 0.3536, 0.3536,
0.5000, 0.3536, 0.5000])
Let's now put all pieces together: We include self-loops, the linear transformation of node features using an instance of torch.nn.Linear, as well as the symmetric degree-based normalization. Conveniently, we can pass a normalization tensor (with one entry per directed edge) to the propagate function. We then apply this during the actual message passing, implemented in the body of message. This is done by reshaping the normalization tensor to shape $(m,1)$ using view(-1,1) and multiplying it element-wise with the node features. Note that in the message function node features are automatically lifted to the level of edges, i.e. x_j has shape $(m,d)$ and includes the $d$-dimensional node features of the source nodes of all (directed) edges. Reshaping norm to shape $(m,1)$ allows it to be broadcast against x_j, so that each message is scaled by the normalization factor of its edge:
class GraphConvolution(torch_geometric.nn.MessagePassing):
"""A graph convolution layer following (Kipf, Welling 2017)."""
# Arguments passed to `propagate()`. pyG usually infers them from the source code of the class,
# which is not possible for classes defined in marimo notebooks
propagate_type = {'x': torch.Tensor, 'norm': torch.Tensor}
def __init__(self, in_ch, out_ch):
"""Initialize the layer parameters."""
super().__init__(aggr='add')
self.linear = torch.nn.Linear(in_ch, out_ch)
def forward(self, x, edge_index):
"""Perform a single round of message passing."""
# we add self loops
edge_index, _ = torch_geometric.utils.add_self_loops(edge_index, num_nodes=x.size(0))
# we linearly transform node features before doing the message passing
x = self.linear(x)
# we normalize features based on the product of the square root of in-degrees
source, target = edge_index
deg = torch_geometric.utils.degree(target, x.size(0), dtype=x.dtype)
deg_inv_sqrt = deg.pow(-0.5)
deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0
norm = deg_inv_sqrt[source] * deg_inv_sqrt[target]
return self.propagate(edge_index, x=x, norm=norm)
def message(self, x_j, norm):
"""Construct the message that is sent along each edge."""
# in our example, x_j has shape [m, 1]
return norm.view(-1, 1) * x_j
This completes the implementation of the Graph Convolution according to Kipf and Welling. We can instantiate our model and calculate the transformed node features after one step of the message passing:
mp_3 = GraphConvolution(1, 1)
_x = mp_3.forward(n.data.x, n.data.edge_index)
print(_x)
tensor([[0.7004],
[0.7004],
[1.3827],
[1.3743],
[1.3743]], grad_fn=<ScatterAddBackward0>)
Note that the model now includes a linear perceptron with two trainable parameters (the weight and bias) that is applied to the feature by each node before passing the (transformed) feature to its neighbors. We can output those parameters, which we will later learn based on Stochastic Gradient Descent:
print([print(x) for x in mp_3.parameters()])
Parameter containing: tensor([[0.3370]], requires_grad=True) Parameter containing: tensor([0.1107], requires_grad=True) [None, None]