Alessio Borgi

Alessio Borgi

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

Molecular GNNs: Learning on Atoms and Bonds

5 minute read

Published:

TL;DR: A molecule is a graph (atoms = nodes, bonds = edges). Molecular GNNs replace hand-crafted fingerprints with learned embeddings. The progression: 2D graphs (connectivity only) to 3D with distances (SchNet) to 3D with angles (DimeNet) to fully equivariant (EGNN, NequIP). Each step adds geometric information and improves accuracy on quantum chemistry benchmarks.

Molecules as Graphs

A molecule G = (V, E, X, R) where:

  • V: atoms (carbon, oxygen, nitrogen, …)
  • E: bonds (single, double, triple, aromatic)
  • X: atom features (atomic number, charge, hybridisation, …)
  • R ∈ ℝ^{N×3}: 3D coordinates (from DFT calculations or conformation search)

The task: predict molecular properties from G. Properties include:

  • HOMO-LUMO gap (electronic structure, relevant to photovoltaics)
  • Solubility (pharmaceutical drug delivery)
  • Toxicity (drug safety screening)
  • Binding affinity (protein-drug interaction)
  • Dipole moment, polarisability (material properties)

From Fingerprints to GNNs

Traditional approach — Morgan fingerprints (ECFP):

  • Encode each atom’s K-hop neighbourhood as a hash
  • Sum over all atoms → fixed-size bit vector
  • Feed to SVM or random forest

GNN approach:

  • Run K rounds of message passing → node embeddings encode K-hop neighbourhoods
  • Global pooling → graph embedding
  • MLP → property prediction

GNNs outperform fingerprints because they learn task-specific features rather than encoding all structural information uniformly.

Level 1: 2D GNNs (Connectivity Only)

Standard GCN, GAT, or GIN on the molecular graph with:

  • Node features: atomic number (one-hot), formal charge, number of Hs, hybridisation
  • Edge features: bond type (single/double/triple/aromatic), is-conjugated, is-ring

Examples: MPNN (Gilmer et al., 2017), AttentiveFP.

Limitation: cannot distinguish stereoisomers (L-alanine vs D-alanine have identical connectivity). Missing 3D information.

Level 2: 3D Distance-Based (Invariant)

Add interatomic distances as edge features. Build a graph where all atoms within cutoff distance r_c are connected (not just bonded atoms).

m_{ij} = φ( h_i, h_j, ||r_i - r_j|| )

SchNet (Schütt et al., 2017): uses continuous-filter convolutions based on radial basis functions of distance:

h_i ← Σ_j h_j · W( ||r_i - r_j|| )

Where W is a distance-dependent filter network. This achieves rotational invariance (distances are invariant) but cannot detect angles.

QM9 performance: SchNet achieves chemical accuracy on several QM9 targets (energy, HOMO energy, LUMO energy) — competitive with DFT at orders-of-magnitude less compute.

Level 3: Angular GNNs (Bond Angles)

Adding interatomic distances is not sufficient — two conformers can have identical pairwise distance matrices but different angles. DimeNet incorporates angles between bond triplets.

DimeNet (Klicpera et al., 2020): messages are defined over directed edges (not just nodes), including the angle between edges:

m_{ji} ← AGG_{k ∈ N(j)} φ( m_{kj}, ||r_{ji}||, θ_{kji} )

Where θ_{kji} is the angle at j between bonds ji and jk. DimeNet uses Bessel functions for radial features and spherical harmonics for angular features.

SphereNet: adds dihedral angles (torsions) — the angle between two planes defined by four atoms. This completes the geometric description of local 3D structure.

Why angles matter: Two carbon atoms bonded to the same central atom at different angles (e.g., 90° vs 120°) experience very different bonding environments. The angle encodes hybridisation (sp³ = 109.5°, sp² = 120°, sp = 180°) and strain. Ignoring angles misses key chemical information.

Level 4: Equivariant GNNs

Equivariant models process 3D positions as vectors, maintaining E(n) or SE(3) equivariance. They can predict both scalar properties (energy) and vector properties (forces) without violating symmetry.

EGNN: distance-based messages + equivariant coordinate updates. Simple, fast, effective for energy prediction.

NequIP: TFN-style tensor features + message passing. State-of-the-art for force fields with few training points.

MACE: many-body interactions. Current SOTA on MD17 molecular dynamics benchmark.

Benchmarks

QM9: 134k small organic molecules (up to 9 heavy atoms). 12 quantum chemical properties (HOMO energy, LUMO energy, dipole moment, etc.) computed by DFT.

MD17: molecular dynamics trajectories. Predict energy and forces at each timestep. Tests generalisation to conformational space.

OGB-molhiv / OGB-molpcba: large-scale drug discovery benchmarks (41k/437k molecules).

PDBbind: protein-ligand binding affinity from crystal structures.

The Accuracy Progression on QM9 (HOMO-LUMO gap)

Morgan fingerprint + RF:  ~0.5 eV
MPNN (2D):                ~0.18 eV
SchNet (distance):        ~0.07 eV
DimeNet (angles):         ~0.05 eV
SphereNet (dihedrals):    ~0.03 eV
NequIP (equivariant):     ~0.02 eV
MACE (many-body):         ~0.01 eV  (near DFT accuracy)

Each geometric level roughly halves the error. Chemical accuracy is ~0.04 eV — equivariant models are within or below this threshold.

Summary

LevelGeometry usedKey modelQM9 error    
2D (connectivity)NoneMPNN, GIN~0.18 eV    
Distances  r_ij  SchNet~0.07 eV
Distances + anglesθ_{ijk}DimeNet~0.05 eV    
Distances + angles + dihedralsφ_{ijkl}SphereNet~0.03 eV    
Full equivariance3D vectorsNequIP, MACE~0.01 eV    

For industrial drug discovery, 2D GNNs suffice for fast virtual screening. For physics-accurate property prediction and force fields, equivariant models are the only option.

References

You May Also Enjoy

Open Problems and Future Directions in Sheaf Neural Networks

5 minute read

Published:

Sheaf Neural Networks are a young field with many open research problems. This post surveys the most important unsolved questions: scalability to million-node graphs, stochastic sheaves for uncertainty quantification, sheaf autoencoders for generative models, connections to persistent homology, and unification with other topological deep learning frameworks.

Sheaf GNNs for Molecular Property Prediction

7 minute read

Published:

Molecular graphs are naturally heterophilic — atoms of different types bonded together — and have rich geometric structure (bond angles, torsion angles). Sheaf GNNs can encode this structure via restriction maps that represent the relational geometry between bonded atoms. This post explores how sheaf theory applies to molecular property prediction.

Sheaf GNNs for Heterophilic Node Classification: Benchmarks and Results

6 minute read

Published:

A systematic comparison of sheaf GNNs against baselines on standard heterophilic node classification benchmarks. Which datasets benefit most from sheaf structure? How does stalk dimension d affect performance? When do sheaves win — and when don’t they?

Sheaves Meet Attention: Transformer-Inspired Sheaf Models

6 minute read

Published:

Graph Transformers use global attention to capture long-range dependencies that local message passing misses. Sheaf Transformers combine sheaf-based local aggregation (with learnable relational geometry) with global attention (for long-range dependencies). This post surveys the design space of sheaf-augmented attention models.