A Gentle Introduction to Geometric Graph Neural Networks

Graph Neural Networks (GNNs) are a part of an emerging research paradigm called Geometric Deep Learning -- devising neural network architectures that respect the invariances and symmetries in data. This practical aims to be a gentle introduction into the world of Geometric Deep Learning.

The aims of this practical are as follows:

• Understanding the concepts of invariance and equivariance to symmetries, which are fundamental properties of Graph Neural Networks. We will cover theory and proofs, as well as programming and unit testing.
• Becoming hands-on with PyTorch Geometric (PyG), a popular libary for developing state-of-the-art GNNs and Geometric Deep Learning models. In particular, gaining familiarity with the MessagePassing base class for designing novel GNN layers and the Data object for representing graph datasets.
• Gaining an appreciation of the fundamental principles behind constructing GNN layers that take advantage of geometric information for graphs embedded in 3D space, such as biomolecules, materials, and other physical systems.

Authors and Acknowledgements

Here are the authors: do not hesitate to reach out to us for any queries and feedback on your solutions!

• Chaitanya K. Joshi (ckj24@cl.cam.ac.uk)
• Charlie Harris (cch57@cam.ac.uk)
• Ramon Viñas Torné (rv340@cam.ac.uk)

This notebook was initially developed for students taking the following courses:

• Representation Learning on Graphs and Networks, at University of Cambridge's Department of Computer Science and Technology (instructors: Prof. Pietro Liò, Dr. Petar Veličković).
• Geometric Deep Learning, at the African Master’s in Machine Intelligence (instructors: Prof. Michael Bronstein, Prof. Joan Bruna, Dr. Taco Cohen, Dr. Petar Veličković).

⚙️ Part 0: Installation and Setup

❗️Note: You will need a GPU to complete this practical. If using Collab, remember to click Runtime -> Change runtime type, and set the hardware accelerator to GPU.

Great! We are ready to dive into the practical!

🧪 Part 0: Introduction to Molecular Property Prediction with PyTorch Geometric

This section covers the fundamentals. We will study how Graph Neural Networks (GNNs) can be employed for predicting chemical properties of molecules, an impactful real-world application of Geometric Deep Learning. To achieve this, we will first introduce PyTorch Geometric, a widely-used Python library that facilitates the implementation of GNNs.

PyTorch Geometric

PyTorch Geometric (PyG) is an excellent library for graph representation learning research and development:

PyTorch Geometric (PyG) consists of various methods for deep learning on graphs and other irregular structures, also known as Geometric Deep Learning, from a variety of published papers. In addition, it provides easy-to-use mini-batch loaders for operating on many small and single giant graphs, multi GPU-support, distributed graph learning, a large number of common benchmark datasets, and helpful transforms, both for learning on arbitrary graphs as well as on 3D meshes or point clouds.

In this practical, we will make extensive use of PyG. If you have never worked with PyG before, do not worry, we will provide you with some examples and guide you through all the fundamentals in a detailed manner. We also highly recommend this self-contained official tutorial, which will help you get started. Among other things, you will learn how to implement state-of-the-art GNN layers via the generic PyG Message Passing class (more on this later).

Now, let's turn our attention to the problem of predicting molecular properties.

Molecular Property Prediction

Molecules are a great example of an object from nature that can easily be represented as a graph of atoms (nodes) connected by bonds (edges). A popular application of GNNs in chemistry is the task of Molecular Property Prediction. The goal is to train a GNN model from historical experimental data that can predict useful properties of drug-like molecules. The model's predictions can then be used to guide the drug design process.

One famous example of GNNs being used in molecular property prediction is in the world of antibiotic discovery, an area with a potentially massive impact on humanity and infamously little innovation. A GNN trained to predict how much a molecule would inhibit a bacteria was able identify the previously overlooked compound Halicin (below) during virtual screening. Not only did halicin show powerful results during in vitro (in cell) testing but it also had a completely novel mechanism of action that no bacteria has developed resistance to (yet).

The QM9 Dataset

QM9 (Quantum Mechanics dataset 9) is a dataset consisting of about 130,000 small molecules with 19 regression targets. Since being used by MoleculeNet, it has become a popular dataset to benchmark new architectures for molecular property prediction.

Specifically, we will be predicting the electric dipole moment of drug-like molecules. According to Wikipedia:

"The electric dipole moment is a measure of the separation of positive and negative electrical charges within a system, that is, a measure of the system's overall polarity."

We can visualize this concept via the water molecule H₂0, which forms a dipole due to its slightly different distribution of negative (blue) and postive (red) charge.

You do not need to worry about the exact physical and chemical principles that underpin dipole moments. As you might imagine, writing the equations from first priciples to predict a property like this, espeically for complex molecules (e.g. proteins), is very difficult. All you need know (for the sake of this practical anyway) is that these molecules can be representated as graphs with node and edge features as well as spatial information that we can use to train a GNN model using the ground truth labels.

Now let us load the QM9 dataset and explore how molecular graphs are represented. PyG makes this extremely convinient.

(The dataset may take a few minutes to download.)

Data Preparation and Splitting

The QM9 dataset has over 130,000 molecular graphs!

Let us create a more tractable sub-set of 3,000 molecular graphs for the purposes of this practical and separate it into training, validation, and test sets. We shall use 1,000 graphs each for training, validation, and testing.

Towards the end of this practical, you will get to experiment with the full/larger sub-sets of the QM9 dataset, too.

Visualising Molecular Graphs

To get a better understanding of how the QM9 molecular graphs look like, let's visualise a few samples from the training set along with their corresponding target (their dipole moment).

In the following plot we visualise sparse graphs where edges represent physical connections (i.e. bonds). In this practical, however, we will use fully-connected graphs and encode the graph structure in the attributes of each. Later in this practical, we will study the advantages and downsides of both approaches.

❗️Note: we have implemented some code for you to convert the PyG graph into a Molecule object that can be used by RDKit, a python package for chemistry and visualing molecules. It is not important for you to understand RDKit beyond visualisation purposes.

Understanding PyG Data Objects

Each graph in our dataset is encapsulated in a PyG Data object, a convient way of representing all structured data for use in Geometric Deep Learning (including graphs, point clouds, and meshes).

Within an instance of a Data object, individual Torch.Tensor attributes (or any other variable type) can be easily dot accessed within a neural network layer. The graphs from PyG come with a number of pre-computed features which we describe below (do not worry if you are unfamiliar with the chemistry terms here):

Atom features (data.x) - ${\mathbb{R}}^{|V|\times 11}$

• 1st-5th features: Atom type (one-hot: H, C, N, O, F)
• 6th feature (also data.z): Atomic number (number of protons).
• 7th feature: Aromatic (binary)
• 8th-10th features: Electron orbital hybridization (one-hot: sp, sp2, sp3)
• 11th feature: Number of hydrogens

Edge Index (data.edge_index) - ${\mathbb{R}}^{2\times |E|}$

• A tensor of dimensions 2 x num_edges that describe the edge connectivity of the graph

Edge features (data.edge_attr) - ${\mathbb{R}}^{|E|\times 4}$

• 1st-4th features: bond type (one-hot: single, double, triple, aromatic)

Atom positions (data.pos) - ${\mathbb{R}}^{|V|\times 3}$

• 3D coordinates of each atom . (We will talk about their importance later in the practical.)

Target (data.y) - ${\mathbb{R}}^1$

• A scalar value corresponding to the molecules electric dipole moment

❗️Note: We will be using fully-connected graphs (i.e. all atoms in a molecule are connected to each other, except self-loops). The information about the molecule structures will be available to the models through the edge features (data.edge_attr) as follows:

• When two atoms are physically connected, the edge attributes indicate the bond type (single, double, triple, or aromatic) through a one-hot vector.
• When two atoms are not physically connected, all edge attributes are zero. We will later study the advantages/downsides of fully-connected adjacency matrices versus sparse adjacency matrices (where an edge between two atoms is present only when there exists a physical connection between them).

Using PyG for batching

As you might remember from the previous practical, batching graphs can be quite a tedious and fiddly process. Thankfully, using PyG makes this super simple! Given a list of Data objects, we can easily batch this into a PyG Batch object as well as unbatch back into a list of graphs. Furthermore, in simple cases like ours, the PyG DataLoader object (different from the vanilla PyTorch one) handles all of the batching under the hood for us!

Lets quicky batch and unbatch some graphs anyway as a demonstration:

Awesome! We have downloaded and prepared the QM9 dataset, visualised some samples, understood the attributes associated with each molecular graph, and reviewed how batching works in PyG. Now, we are ready to understand how we can develop GNNs in PyG for molecular property prediction.

📩 Part 0: Introduction to Message Passing Neural Networks in PyTorch Geometric

As a gentle introduction to PyTorch Geometric, we will walk you through the first steps of developing a GNN in the Message Passing flavor.

Formalism

Firstly, let us formalise our molecular property prediction pipeline. (Our notation will mostly follow what has been introduced in the lectures, but we do make some different choices for variable names.)

Graph

Consider a molecular graph $\mathcal{G}=\left(\mathcal{V},\mathcal{E}\right)$, where $\mathcal{V}$ is a set of $n$ nodes, and $\mathcal{E}$ is a set of edges associated with the nodes. For each node $i\in \mathcal{V}$, we are given a $d_n$-dimensional initial feature vector $h_i\in {\mathbb{R}}^{d_n}$. For each edge $(i,j)\in \mathcal{E}$, we are given a $d_e$-dimensional initial feature vector $e_{ij}\in {\mathbb{R}}^{d_e}$. For QM9 graphs, $d_n=11,d_e=4$.

Label/target

Associated with each graph $\mathcal{G}$ is a scalar target or label $y\in {\mathbb{R}}^1$, which we would like to predict.

We will design a Message Passing Neural Network for graph property prediction to do this. Our MPNN will consist of several layers of message passing, followed by a global pooling and prediction head.

MPNN Layer

The Message Passing operation iteratively updates node features $h_i^{\ell }\in {\mathbb{R}}^d$ from layer $\ell$ to layer $\ell +1$ via the following equation: $$h_i^{\ell +1}=\varphi (h_i^{\ell },{\oplus }_{j\in {\mathcal{N}}_i}(\psi \left(h_i^{\ell },h_j^{\ell },e_{ij}\right)\left)\right),$$ where $\psi ,\varphi$ are Multi-Layer Perceptrons (MLPs), and $\oplus$ is a permutation-invariant local neighborhood aggregation function such as summation, maximization, or averaging.

Let us break down the MPNN layer into three pedagogical steps:

• Step (1): Message. For each pair of linked nodes $i,j$, the network first computes a message $m_{ij}=\psi \left(h_i^{\ell },h_j^{\ell },e_{ij}\right)$. The MLP $\psi :{\mathbb{R}}^{2d+d_e}\to {\mathbb{R}}^d$ takes as input the concatenation of the feature vectors from the source node, destination node, and edge.
  • Note that for the first layer $\ell =0$, $h_i^{\ell =0}=W_{in}\left(h_i\right)$, where $W_{in}\in {\mathbb{R}}^{d_n}\to {\mathbb{R}}^d$ is a simple linear projection (torch.nn.Linear) for the initial node features to hidden dimension $d$.
• Step (2): Aggregate. At each node $i$, the incoming messages from all its neighbors are then aggregated as $m_i={\oplus }_{j\in {\mathcal{N}}_i}\left(m_{ij}\right)$, where $\oplus$ is a permutation-invariant function. We will use summation, i.e. ${\oplus }_{j\in {\mathcal{N}}_i}={\sum }_{j\in {\mathcal{N}}_i}$.
• Step (3): Update. Finally, the network updates the node feature vector $h_i^{\ell +1}=\varphi \left(h_i^{\ell },m_i\right)$, by concatenating the aggregated message $m_i$ and the previous node feature vector $h_i^{\ell }$, and passing them through an MLP $\varphi :{\mathbb{R}}^{2d}\to {\mathbb{R}}^d$.

Global Pooling and Prediction Head

After $L$ layers of message passing, we obtain the final node features $h_i^{\ell =L}$. As we have a single target $y$ per graph, we must pool all node features into a single graph feature or graph embedding $h_G\in {\mathbb{R}}^d$ via another permutation-invariant function $R$, sometimes called the 'readout' function, as follows: $$h_G=R_{i\in \mathcal{V}}\left(h_i^{\ell =L}\right).$$ We will use global average pooling over all node features, i.e. $$h_G=\frac{1}{|\mathcal{V}|}\sum _{i\in \mathcal{V}}{h}_i^{\ell =L}.$$

The graph embedding $h_G$ is passed through a linear prediction head $W_{pred}\in {\mathbb{R}}^d\to {\mathbb{R}}^1$ to obtain the overall prediction $\hat{y}\in {\mathbb{R}}^1$: $$\hat{y}=W_{pred}\left(h_G\right).$$

Loss Function

Our MPNN graph property prediction model can be trained end-to-end via minimizing the standard mean-squared error loss for regression: $${\mathcal{L}}_{MSE}=\Vert y-\hat{y}{\Vert }_2^2.$$

Coding the basic Message Passing Neural Network Layer

We are now ready to define a basic MPNN layer which implements what we have described above. In particular, we will code up the MPNN Layer first. (We will code up the other parts subsequently.)

To do so, we will inherit from the MessagePassing base class, which automatically takes care of message propagation and is extremely useful to develop advanced GNN models. To implement a custom MPNN, the user only needs to define the behaviour of the message (i.e. $\psi$), the aggregate(i.e. $\oplus$), and update (i.e. $\varphi$) functions. You may also refer to the PyG documentation for implementing custom message passing layers.

Below, we provide the implementation of a standard MPNN layer as an example, with extensive inline comments to help you figure out what is going on.

Great! We have defined a Message Passing layer following the equation we had introduced previously. Let us use this layer to code up the full MPNN graph property prediction model. This model will take as input molecular graphs, process them via multiple MPNN layers, and predict a single property for each of them.

Awesome! We are done defining our first MPNN model for graph property prediction.

But wait! Before we dive into training and evaluation this model, let us write some sanity checks for a fundamental property of the model and the layer.

Unit tests for Permutation Invariance and Equivariance

The lectures have repeatedly indicated on certain fundamental properties for machine learning on graphs:

• A GNN layer is equivariant to permutations of the set of nodes in the graph; i.e. as we permute the nodes, the node features produced by the GNN must permute accordingly.
• A GNN model for graph-level property prediction is invariant to the permutations of the set of nodes in the graph; i.e. as we permute the nodes, the graph-level propery remains unchanged.

(But wait...What is a permutation? Essentially, it is an ordering of the nodes in a graph. In general, there is no canonical way of assigning an ordering of the nodes, unlike textual or image data. However, graphs need to be stored and processed on computers in order to perform machine learning on them (which is what this course is about!). Thus, we need to ensure that our models are able to principaly handle this lack of canonical ordering or permutation of graph nodes. This is what the above statements are trying to say.)

Formalism

Let us try to formalise these notions of permutation invariance and equivariance via matrix notation (it is easier that way).

• Let $\mathbf{H}\in {\mathbb{R}}^{n\times d}$ be a matrix of node features for a given molecular graph, where $n$ is the number of nodes/atoms and each row $h_i$ is the $d$-dimensional feature for node $i$.
• Let $\mathbf{A}\in {\mathbb{R}}^{n\times n}$ be the adjacency matrix where each entry denotes $a_{ij}$ the presence or absence of an edge between nodes $i$ and $j$.
• Let $\mathbf{F}(\mathbf{H},\mathbf{A}):{\mathbb{R}}^{n\times d}\times {\mathbb{R}}^{n\times n}\to {\mathbb{R}}^{n\times d}$ be a GNN layer that takes as input the node features and adjacency matrix, and returns the updated node features.
• Let $f(\mathbf{H},\mathbf{A}):{\mathbb{R}}^{n\times d}\times {\mathbb{R}}^{n\times n}\to \mathbb{R}$ be a GNN model that takes as input the node features and adjacency matrix, and returns the predicted graph-level property.
• Let $\mathbf{P}\in {\mathbb{R}}^{n\times n}$ be a permutation matrix which has exactly one 1 in every row and column, and 0s elsewhere. Left-multipying $\mathbf{P}$ with a matrix changes the ordering of the rows of the matrix.

Permuation Equivariance

The GNN layer $\mathbf{F}$ is permuation equivariant as follows: $$\mathbf{F}(\mathbf{PH},\mathbf{PA}{\mathbf{P}}^{\mathbf{T}})=\mathbf{P}\mathbf{F}(\mathbf{H},\mathbf{A}).$$

Another way to formulate the above could be: (1) Consider the updated node features ${\mathbf{H}}^{\prime }=\mathbf{F}(\mathbf{H},\mathbf{A})$. (2) Applying any permutation matrix $\mathbf{P}$ to the input of the GNN layer $\mathbf{F}$ should produce the same result as applying the same permutation on ${\mathbf{H}}^{\prime }$: $$\mathbf{F}(\mathbf{PH},\mathbf{PA}{\mathbf{P}}^{\mathbf{T}})=\mathbf{P}{\mathbf{H}}^{\prime }$$

Permuation Invariance

The GNN model $f$ for graph-level prediction is permutation invariant as follows: $$f(\mathbf{PH},\mathbf{PA}{\mathbf{P}}^{\mathbf{T}})=f(\mathbf{H},\mathbf{A}).$$

Another way to formulate the above could be: (1) Consider the predicted molecular property $\hat{\mathbf{y}}=f(\mathbf{H},\mathbf{A})$. (2) Applying any permutation matrix $\mathbf{P}$ to the input of the GNN model $f$ should produce the same result as not applying it: $$f(\mathbf{PH},\mathbf{PA}{\mathbf{P}}^{\mathbf{T}})=\hat{\mathbf{y}}.$$

With that formalism out of the way, let us write some unit tests to confirm that our MPNNModel and MPNNLayer are indeed permutation invariant and equivariant, respectively.

Now that we have defined the unit tests for permutation invariance (for the full MPNN model) and permutation equivariance (for the MPNN layer), let us perform the sanity check:

Training and Evaluating Models

Great! We are finally ready to train and evaluate our model on QM9. We have provided a basic experiment loop which takes as input the model and dataloaders, performs training, and returns the final performance on the validation and test set.

We will be training a MPNNModel consisting of 4 layers of message passing with a hidden dimension of 64.

Super! Everything up to this point has already been covered in the lectures, and we hope that the practical so far has been a useful recap along with the acompanying code.

Now for the fun part, where you will be required to think what you have studied so far!

🧊 Part 1: Geometric Graphs and Message Passing with 3D Coordinates

Remember that we were given 3D coordinates with each atom in our molecular graph?

Molecular graphs, and other structured data occurring in nature, do not simply exist on flat planes. Instead, molecules have an inherent 3D structure that influences their properties and functions.

Let us visualize a molecule from QM9 in all of its 3D glory!

Go ahead and try move this molecule with your mouse cursor!

💻Task 1.1: Develop a Message Passing Neural Network that incorporates the atom coordinates as node features (0.5 Marks).

Our initial and somewhat 'vanilla' MPNN MPNNModel ignores the atom coordiantes and only uses the node features to perform message passing. This means that the model is not leveraging useful 3D structural information to predict the target property.

Your first task is to modify the original MPNNModel to incorporate atom coordinates into the node features.

We have defined most of the new CoordMPNNModel class for you, and you have to fill in the YOUR CODE HERE sections.

🤔 Hint: As reminder, the 3D atom positions are stored in data.pos. You don't have to do something very smart right now (that will come later). A simple solution is okay to get started, e.g. concatenation or summation.

💻Task 1.2: Test the permutation invariance and equivariance properties of your new CoordMPNNModel with node features and coordinates, as well as the constituent MPNNLayer. (0.5 Marks)

Super! You have successfully implemented an MPNN which utilises both the atom features as well as coordinates to predict molecular properties.

Before we evaluate it, let us once again run the permutation sanity checks again to make sure the model and layer have the desired properties that constitute every basic GNN:

• The MPNNLayer should be permutation equivariant (we have already shown this previously, but we want you to repeat the exercise in order to thoroughly understand it).
• The CoordMPNNModel should be permutation invariant.

Your task is to fill in the YOUR CODE HERE sections to run the required unit tests. You do not need to write new unit tests yet, the ones we defined previously can be re-used.

💻Task 1.3. Prove that your new CoordMPNNModel is invariant to permutations of both the node features as well as node coordinates. (0.5 Marks)

🤔 Hint: We are looking for simple statements that follow how we formalised permuation invariance for the vanilla MPNN model. We expect you to be copy-pasting most of the formalism and accounting for how your MPNN incorporates both the node features and coordinates. You can additionally introduce $\mathbf{X}\in {\mathbb{R}}^{n\times 3}$ as the matrix of node coordinates for a given molecular graph.

❗️YOUR ANSWER HERE

💻Task 1.4. Train and evaluate your CoordMPNNModel with node features and coordinates on QM9. (0.5 Marks)

Awesome! You are now ready to train and evaluate our new MPNN with node features and coordinates on QM9.

Re-use the experiment loop we have provided and fill in the YOUR CODE HERE sections to run the experiment.

You will be training a CoordMPNNModel consisting of 4 layers of message passing with a hidden dimension of 64, in order to compare your result fairly to the previous vanilla MPNNModel.

Hmm... If you've implemented the CoordMPNNModel correctly up till now, you may see a very curious result -- the performance of CoordMPNNModel is about equal or marginally worse than the vanilla MPNNModel!

This is because the CoordMPNNModel is not using 3D structural information in a principled manner.

The next sections will help us formalise and understand why this is happening.

🔄 Part 2: Invariance to 3D Symmetries: Rotation and Translation

We saw that the performance of CoordMPNNModel is unexpectedly mediocre compared to MPNNModel despite using both node features and coordinates. (But please do not panic if your results say otherwise.) In order to determine why, we must understand the concept of 3D symmetries.

Geometric Invariance

Recall that molecular graphs have 3D coordinates for each atom. A key detail which we have purposely withheld from you up till this point (😈) is that these 3D coordinates are not inherently fixed or permanent. Instead, they were experimentally determined relative to a frame of reference.

To fully grasp these statements, here is GIF of a drug-like molecules moving around in 3D space...

The atoms' 3D coordinates are constantly rotating and translating. However, the properties of this molecule will always remain the same no matter how we rotate or translate it. In other words, the molecule's properties are invariant to 3D rotations and translations.

In this block we will study how to design GNN layers and models that respect these regularities.

Formalism

Let us try to formalise the notion of invariance to 3D rotations and translations in GNNs via matrix notation.

• Let $\mathbf{H}\in {\mathbb{R}}^{n\times d}$ be a matrix of node features for a given molecular graph, where $n$ is the number of nodes/atoms and each row $h_i$ is the $d$-dimensional feature for node $i$.
• Let $\mathbf{X}\in {\mathbb{R}}^{n\times 3}$ be a matrix of node coordinates for a given molecular graph, where $n$ is the number of nodes/atoms and each row $x_i$ is the 3D coordinate for node $i$.
• Let $\mathbf{A}\in {\mathbb{R}}^{n\times n}$ be the adjacency matrix where each entry denotes $a_{ij}$ the presence or absence of an edge between nodes $i$ and $j$.
• Let $\mathbf{F}(\mathbf{H},\mathbf{X},\mathbf{A}):{\mathbb{R}}^{n\times d}\times {\mathbb{R}}^{n\times 3}\times {\mathbb{R}}^{n\times n}\to {\mathbb{R}}^{n\times d}$ be a GNN layer that takes as input the node features, node coordinates, and adjacency matrix, and returns the updated node features.
• Let $f(\mathbf{H},\mathbf{X},\mathbf{A}):{\mathbb{R}}^{n\times d}\times {\mathbb{R}}^{n\times 3}\times {\mathbb{R}}^{n\times n}\to \mathbb{R}$ be a GNN model that takes as input the node features, node coordinates, and adjacency matrix, and returns the predicted graph-level property.

(Notice that we have updated the notation for the GNN layer $\mathbf{F}$ and GNN model $\mathbf{f}$ to include the matrix of node coordinates $\mathbf{X}$ as an additional input.)

💻Task 2.1: What does it mean for the GNN model $f$ and the GNN layer $\mathbf{F}$ to be invariant to 3D rotations and translations? Express this mathematically using the definitions above. (0.5 Mark)

🤔 Hint: Revisit the formalisms for permutation invariance and equivariance to get an idea of how to go about this. You should use the matrix notation we have provided above. Similar to the permuatation matrix $\mathbf{P}$, you may now define an orthogonal rotation matrix $\mathbf{Q}\in {\mathbb{R}}^{3\times 3}$ and a translation vector $\mathbf{t}\in {\mathbb{R}}^3$ in your answer. These would operate on the matrix of node coordinates $\mathbf{X}\in {\mathbb{R}}^{n\times 3}$.

❗️YOUR ANSWER HERE

Before you start coding up a more principled MPNN model, we would like you to take a moment to think about why invariance to 3D rotations and translations is something desirable for GNNs predicting molecular properties...

💻Task 2.2: Is invariance to 3D rotations and translations a desirable property for GNNs? Explain why. (0.5 Marks)

🤔 Hint: We are not looking for an essay, a few sentences will suffice here.

❗️YOUR ANSWER HERE

💻Task 2.3: Write the unit test to check your CoordMPNNModel for 3D rotation and translation invariance. (0.5 Mark)

🤔 Hint: Show that the output of the model varies when:

1. All the atom coordinates in data.pos are multiplied by any random orthogonal rotation matrix $Q\in {\mathbb{R}}^{3\times 3}$. (We have provided a helper function for creating rotation matrices.)
2. All the atom coordinates in data.pos are displaced by any random translation vector $\mathbf{t}\in {\mathbb{R}}^3$.

Now that you have defined the unit tests for rotation and translation invariance, perform the sanity check on your CoordMPNNModel:

(Spoiler alert: if you have implemented things as expected, the unit test should return False for the CoordMPNNModel.)

In this part, you have formalised how a GNN can be 3D rotation and translation invariant, thought about why this is desirable for molecular property prediction, and shown that the CoordMPNNModel was not rotation and translation invariant.

At this point, you should have a concrete understanding of why the performance of CoordMPNNModel is equal or worse than the vanilla MPNNModel, and what we meant by our initial statement before we began this part:

"The CoordMPNNModel is not using 3D structural information in a principled manner"

Let us try fixing this in the next part!

✈️ Part 3: Message Passing with Invariance to 3D Rotations and Translations

This section will dive into how we may design GNN models which operate on graphs with 3D coordinates in a more theoretically sound way.

💻Task 3.1: Design a new Message Passing Layer as well as the accompanying MPNN Model that are both invariant to 3D rotations and translations. (2 Marks)

❗️ Note: There is no single correct answer to this question.

Our initial 'vanilla' MPNN MPNNModel and MPNNLayer ignored the atom coordiantes and only uses the node features to perform message passing. This means that the model was not leveraging 3D structural information to predict the target property.

Our second 'naive' coordinate MPNN CoordMPNNModel used the node features along with the atom coordinates in an unprincipled manner, resulting in the model not being invariant to 3D rotations and translations of the coordinates (which was a desirable property, as we saw in the previous part).

Your task is to define a new InvariantMPNNLayer which utilise both atom coordinates and node features.

We have defined most of the new InvariantMPNNLayer, and you have to fill in the YOUR CODE HERE sections. We have also already defined the InvariantMPNNModel that instantiates your new layer to compose the model. You only need to define the new layer.

🤔 Hint 1: Unlike the previous CoordMPNNModel, we would suggest using the coodinate information to constuct the messages as opposed to incorporating it into the node features. In particular, we would like you to think about how to use the coordinates in a principled manner to constuct the messages: What is a measurement that we can computer using a pair of coordinates that will be invariant to rotating and translating them?

🤔 Hint 2: tensors passed to propagate() can be mapped to the respective nodes and by appending _i or _j to the variable name, e.g. h_i and h_j for the node features h. Note that we generally refer to _i as the central nodes that aggregates information, and refer to _j as the neighboring nodes.