GNN Molecular

Seven GNNs for Molecular Properties

Research · Oct 2024 · Completed

Comparing seven graph neural networks at predicting how well molecules dissolve.

ESOL RMSE 0.498
R² 0.926
7 GNNs compared

Built withPyTorchPyTorch GeometricRDKitMoleculeNet (ESOL)pytest

Code is in a private repository for now.

Overview

A molecule is naturally a graph: atoms are nodes and bonds are edges. I wanted to see how different graph neural network designs use that structure to predict a real chemical property, aqueous solubility, from the MoleculeNet ESOL benchmark.

I first implemented three architectures by hand in notebooks: a graph convolutional network (GCN), a message-passing neural network (MPNN) that also uses bond features, and a diffusion-convolutional network (DCNN) that spreads information over several hops. I later turned them into a tested Python package and added four standard models built from PyTorch Geometric layers: GIN, GAT (v2), GraphSAGE and AttentiveFP.

All seven run on the same seeded split with the same metrics. AttentiveFP (test RMSE 0.498, R² 0.926) and GIN (0.528, 0.916) lead, followed by the MPNN (0.645), GAT (0.769), GraphSAGE (0.816), DCNN (0.907) and GCN (0.981). The four models that read bond types come out ahead of the three that only see atoms.

The question

How do seven graph neural network designs compare at predicting molecular solubility (ESOL) directly from molecular graphs?

Key features

Seven GNNs for graph-level regression, sharing one interface: GCN, MPNN, DCNN, GIN, GAT, GraphSAGE and AttentiveFP.
SMILES-to-graph featurization with RDKit: 9 atom features and 3 bond features.
Bond-aware models (MPNN, GIN, GAT, AttentiveFP) next to atom-only models (GCN, DCNN, GraphSAGE) for a direct comparison.
Two commands: mgnn-train trains any model and reports test metrics, mgnn-predict scores any SMILES string.
75 pytest tests, including permutation invariance and equivariance and batch independence, run on every model.
Seeded 80/10/10 split with MAE, MSE, RMSE, and R² tracked every epoch, plus early stopping on validation loss.

How it works

Molecules to graphs
ESOL molecules are loaded from MoleculeNet and represented as graphs of atoms and bonds, with molecules that contain isolated atoms filtered out where needed.
Batch the graphs
PyTorch Geometric merges many small molecular graphs into one disconnected graph per batch and keeps track of which atom belongs to which molecule.
Propagate information
Each architecture updates atom representations in its own way: graph convolution, learned messages over bonds, multi-hop diffusion, sum-and-MLP (GIN), attention over neighbours (GAT, AttentiveFP) or separate self and neighbour weights (GraphSAGE).
Pool and predict
Atom representations are pooled into one vector per molecule, which a regression head maps to a solubility value, trained with mean squared error.

Examples

AttentiveFP, the best of the seven: the best-validation checkpoint (epoch 222) scores test RMSE 0.498 and R² 0.926.

GIN with bond features: test RMSE 0.528 and R² 0.916, close to AttentiveFP with 25 times fewer parameters.

MPNN training (left) and validation (right) metrics on ESOL. The best-validation checkpoint (epoch 195) scores test RMSE 0.645 and R² 0.875.

The first ESOL molecule (amygdalin), rendered with RDKit: 32 atoms become 32 graph nodes with 9 features each.

Input

mgnn-predict --checkpoint checkpoints/mpnn/best.pt "CCO" "c1ccccc1"

Output

1.195 log(mol/L) CCO -1.633 log(mol/L) c1ccccc1

Real output of the trained MPNN for ethanol and benzene. Their measured values are 1.10 and -1.64; both molecules are part of ESOL, so this shows the tool working, not generalisation.

Tech stack

PyTorch
Model definitions and training loops
PyTorch Geometric
Graph layers (GCNConv, GINEConv, GATv2Conv, SAGEConv, AttentiveFP), message passing, and batching
RDKit
Parse SMILES and inspect molecules
MoleculeNet (ESOL)
Molecular solubility benchmark
pytest
Unit tests for the models, data pipeline, and training loop

What I learned

Bond features matter: every model that reads bond types beats every model that does not
Sum beats mean: GIN's sum aggregation can count neighbours, so it clearly beats the mean-based GraphSAGE and GCN
Tests catch real bugs: one found a DCNN mean pooled over padding, and an eval-mode check showed GraphSAGE had learned to depend on dropout

Limitations

Single dataset (ESOL) and a single seed; private repository

Keywords

graph neural networks · GNN · GCN · graph convolutional network · MPNN · message passing neural network · DCNN · AttentiveFP · GraphSAGE · graph attention network · GAT · graph isomorphism network · GIN · diffusion convolutional neural network · PyTorch Geometric · PyTorch · RDKit · cheminformatics · molecular property prediction · molecular graphs · SMILES · MoleculeNet · ESOL · aqueous solubility prediction · drug discovery · deep learning