Graph Neural Networks for Molecular Data

Molecules are inherently graph-structured objects: atoms are nodes, bonds are edges. Unlike tabular data where all samples share the same feature space, molecular graphs are variable-size, irregular structures where the spatial arrangement of atoms determines chemical and biological properties. Graph Neural Networks (GNNs) exploit this natural graph structure, learning molecular representations end-to-end from the atomic and bond-level features, without the information loss introduced by flattening a molecule into a fixed-length fingerprint vector.

This guide covers the theoretical foundation of GNNs for molecular data, key architectures, node and edge feature engineering, practical implementation with PyTorch Geometric, graph batching mechanics, and a careful analysis of when GNNs outperform fingerprint-based models.

Why Molecules Are Naturally Graphs

A molecule is a weighted undirected graph G = (V, E) where:

• V (nodes/vertices): each heavy atom in the molecule
• E (edges): each bond between two atoms
• Node features h_v: per-atom attributes (atomic number, hybridization, charge, etc.)
• Edge features e_uv: per-bond attributes (bond type, aromaticity, etc.)

This representation is complete: it captures all chemical information present in the connectivity (2D topology). For 3D-sensitive properties (binding pose, conformational effects), the representation can be extended with 3D coordinates as additional node or edge features.

The key advantage over fingerprints is that GNNs learn to compute representations of atom neighborhoods in a task-specific way. Where Morgan fingerprints use fixed-radius circular summaries with no awareness of what property matters, a GNN learns which structural patterns to aggregate based on the training signal from the property label.

Message Passing Neural Networks (MPNN) Intuition

The MPNN framework (Gilmer et al. 2017) unifies most GNN architectures for molecular property prediction into a common formalism with three components:

1. Message: Each atom v sends a message to its neighbor u, computed from v's current representation and the edge feature e_uv: m_v = M(h_v, e_uv)
2. Aggregation: Each atom u collects all incoming messages and aggregates them: a_u = SUM over neighbors v of m_v (or mean, max)
3. Update: Each atom updates its representation using the aggregated message: h_u' = U(h_u, a_u) (often a GRU or MLP)

This is repeated for T message-passing steps (typically 3-6). After T steps, each atom's representation encodes information about its T-hop neighborhood. A final readout function (sum, mean, or attention-weighted sum over all atom representations) produces the molecular representation used for property prediction.

Conceptually, MPNN at T steps is analogous to Morgan fingerprints at radius T, except that the aggregation and encoding functions are learned rather than fixed by algorithm.

Node and Edge Features in Molecular GNNs

Atom (Node) Features

Bond (Edge) Features

A typical atom feature vector has 70-100 dimensions. In 3D GNNs (SchNet, DimeNet), edge features are replaced by interatomic distances and angles, requiring conformer geometry as input.

Key GNN Architectures

Libraries: PyG and DGL

Two main Python libraries implement GNNs for graphs including molecules:

• PyTorch Geometric (PyG): Tightly integrated with PyTorch. Uses sparse adjacency representations via edge_index (2 x num_edges COO format). Provides message-passing base class, many built-in layers (GCN, GAT, MPNN, etc.), and molecular featurizers via torch_geometric.utils.smiles or custom featurizers.
• DGL (Deep Graph Library): Framework-agnostic (PyTorch, MXNet). Has its own message-passing API. DGL-LifeSci provides ready-to-use molecular graph featurizers and model implementations (AttentiveFP, MPNN, etc.).

Code Example: Molecular GNN with PyTorch Geometric

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.data import Data, DataLoader
from torch_geometric.nn import MessagePassing, global_mean_pool, global_add_pool
from torch_geometric.utils import add_self_loops
from rdkit import Chem
from rdkit.Chem import AllChem
import numpy as np
from typing import List, Optional

# -------------------------------------------------------
# 1. Molecule to PyG Data object
# -------------------------------------------------------
ATOM_TYPES = ['C', 'N', 'O', 'S', 'F', 'Cl', 'Br', 'I', 'P', 'Si', 'Other']
HYBRIDIZATION = [Chem.rdchem.HybridizationType.SP,
                 Chem.rdchem.HybridizationType.SP2,
                 Chem.rdchem.HybridizationType.SP3,
                 Chem.rdchem.HybridizationType.SP3D,
                 Chem.rdchem.HybridizationType.SP3D2]
BOND_TYPES = [Chem.rdchem.BondType.SINGLE,
              Chem.rdchem.BondType.DOUBLE,
              Chem.rdchem.BondType.TRIPLE,
              Chem.rdchem.BondType.AROMATIC]

def one_hot(value, choices):
    return [int(value == c) for c in choices] + [int(value not in choices)]

def atom_features(atom) -> List[int]:
    """Compute 72-dimensional atom feature vector."""
    symbol = atom.GetSymbol()
    features = (
        one_hot(symbol, ATOM_TYPES) +                           # 12
        one_hot(atom.GetDegree(), list(range(11))) +            # 12
        one_hot(atom.GetHybridization(), HYBRIDIZATION) +       # 6
        [int(atom.GetIsAromatic())] +                           # 1
        one_hot(atom.GetFormalCharge(), [-2,-1,0,1,2]) +        # 6
        one_hot(atom.GetTotalNumHs(), [0,1,2,3,4]) +            # 6
        [int(atom.IsInRing())] +                                # 1
        one_hot(atom.GetChiralTag(), [
            Chem.rdchem.ChiralType.CHI_UNSPECIFIED,
            Chem.rdchem.ChiralType.CHI_TETRAHEDRAL_CW,
            Chem.rdchem.ChiralType.CHI_TETRAHEDRAL_CCW,
            Chem.rdchem.ChiralType.CHI_OTHER,
        ])                                                       # 5
    )
    return features  # total: 12+12+6+1+6+6+1+5 = 49 dims

def bond_features(bond) -> List[int]:
    """Compute 6-dimensional bond feature vector."""
    return (
        one_hot(bond.GetBondType(), BOND_TYPES) +    # 5
        [int(bond.GetIsConjugated())] +              # 1
        [int(bond.IsInRing())]                       # 1
    )  # total: 7 dims

def mol_to_graph(smiles: str, label: Optional[float] = None) -> Optional[Data]:
    """Convert SMILES to a PyG Data object."""
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        return None

    # Node features
    x = torch.tensor([atom_features(a) for a in mol.GetAtoms()], dtype=torch.float)

    # Edge index and edge features
    rows, cols, edge_attrs = [], [], []
    for bond in mol.GetBonds():
        i = bond.GetBeginAtomIdx()
        j = bond.GetEndAtomIdx()
        bf = bond_features(bond)
        rows += [i, j]
        cols += [j, i]
        edge_attrs += [bf, bf]  # Undirected: add both directions

    edge_index = torch.tensor([rows, cols], dtype=torch.long)
    edge_attr  = torch.tensor(edge_attrs, dtype=torch.float)

    y = torch.tensor([label], dtype=torch.float) if label is not None else None
    return Data(x=x, edge_index=edge_index, edge_attr=edge_attr, y=y)

# -------------------------------------------------------
# 2. MPNN Layer
# -------------------------------------------------------
class MPNNConv(MessagePassing):
    """
    One MPNN message-passing step with edge feature integration.
    Uses 'add' aggregation and an MLP for message computation.
    """
    def __init__(self, node_dim: int, edge_dim: int, out_dim: int):
        super().__init__(aggr='add')
        # Message network: maps (node_src, edge) -> message
        self.msg_mlp = nn.Sequential(
            nn.Linear(node_dim + edge_dim, out_dim),
            nn.ReLU(),
            nn.Linear(out_dim, out_dim),
        )
        # Update network: maps (node + aggregated message) -> new node
        self.update_gru = nn.GRUCell(out_dim, node_dim)

    def forward(self, x, edge_index, edge_attr):
        return self.propagate(edge_index, x=x, edge_attr=edge_attr)

    def message(self, x_j, edge_attr):
        # x_j: source node features, edge_attr: bond features
        return self.msg_mlp(torch.cat([x_j, edge_attr], dim=-1))

    def update(self, aggr_out, x):
        # GRU update: combine current node state with aggregated messages
        return self.update_gru(aggr_out, x)

# -------------------------------------------------------
# 3. Full MPNN Molecular Model
# -------------------------------------------------------
class MolMPNN(nn.Module):
    def __init__(self, node_in_dim: int = 49, edge_in_dim: int = 7,
                 hidden_dim: int = 256, n_layers: int = 4,
                 n_tasks: int = 1, dropout: float = 0.1):
        super().__init__()

        # Input projection
        self.node_proj = nn.Linear(node_in_dim, hidden_dim)
        self.edge_proj = nn.Linear(edge_in_dim, hidden_dim)

        # Message passing layers
        self.mp_layers = nn.ModuleList([
            MPNNConv(hidden_dim, hidden_dim, hidden_dim)
            for _ in range(n_layers)
        ])

        # Readout: sum over all atoms
        self.dropout = nn.Dropout(dropout)

        # Output MLP
        self.mlp = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim // 2, n_tasks),
        )

    def forward(self, data):
        x, edge_index, edge_attr, batch = (
            data.x, data.edge_index, data.edge_attr, data.batch
        )

        # Initial projections
        x = F.relu(self.node_proj(x))
        edge_attr = F.relu(self.edge_proj(edge_attr))

        # Message passing
        for mp_layer in self.mp_layers:
            x = x + self.dropout(mp_layer(x, edge_index, edge_attr))  # Residual

        # Global pooling: sum all atom representations
        mol_repr = global_add_pool(x, batch)  # (batch_size, hidden_dim)

        return self.mlp(mol_repr)  # (batch_size, n_tasks)

# -------------------------------------------------------
# 4. Training on ESOL (aqueous solubility) dataset
# -------------------------------------------------------
from torch_geometric.datasets import MoleculeNet

# Load ESOL dataset (1128 molecules, regression task)
dataset = MoleculeNet(root='/tmp/esol', name='ESOL')

# Convert to train/val/test split
n = len(dataset)
train_data = dataset[:int(0.8 * n)]
val_data   = dataset[int(0.8 * n):int(0.9 * n)]
test_data  = dataset[int(0.9 * n):]

train_loader = DataLoader(train_data, batch_size=32, shuffle=True)
val_loader   = DataLoader(val_data, batch_size=32)
test_loader  = DataLoader(test_data, batch_size=32)

# Infer node/edge dims from dataset
sample = dataset[0]
node_dim = sample.x.shape[1]   # 9 in MoleculeNet ESOL
edge_dim = sample.edge_attr.shape[1] if sample.edge_attr is not None else 0

model = MolMPNN(
    node_in_dim=node_dim,
    edge_in_dim=max(edge_dim, 1),
    hidden_dim=128,
    n_layers=4,
    n_tasks=1,
    dropout=0.1
)

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.MSELoss()

def train_epoch(loader):
    model.train()
    total_loss = 0.0
    for data in loader:
        out = model(data).squeeze(-1)
        y   = data.y.squeeze(-1)
        loss = criterion(out, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * data.num_graphs
    return total_loss / len(loader.dataset)

def evaluate(loader):
    model.eval()
    preds, targets = [], []
    with torch.no_grad():
        for data in loader:
            preds.append(model(data).squeeze(-1))
            targets.append(data.y.squeeze(-1))
    preds   = torch.cat(preds)
    targets = torch.cat(targets)
    rmse = ((preds - targets) ** 2).mean().sqrt().item()
    return rmse

for epoch in range(1, 101):
    train_loss = train_epoch(train_loader)
    val_rmse   = evaluate(val_loader)
    if epoch % 10 == 0:
        print(f"Epoch {epoch:3d}: train_loss={train_loss:.4f}, val_RMSE={val_rmse:.4f}")

test_rmse = evaluate(test_loader)
print(f"Test RMSE: {test_rmse:.4f}")

When GNNs Outperform Fingerprints

The question of when to use a GNN versus ECFP4 + gradient boosting is practical and depends on several factors:

The practical recommendation: always establish a fingerprint baseline first. If it achieves adequate performance, stop there. Invest GNN development effort only when the baseline shows a clear ceiling (scaffold split evaluation significantly worse than random split) or when the property has a known dependence on 3D structure.

Practical Tips: Batching and Variable Graph Sizes

Unlike tabular data where all samples have the same feature dimension, molecular graphs have variable numbers of atoms and bonds. PyG handles this through a super-graph batching mechanism: multiple molecular graphs are merged into a single disconnected graph, with a batch index tensor mapping each node to its originating molecule. This allows vectorized GPU computation across the batch without padding.

• The global_add_pool and global_mean_pool functions use the batch index to aggregate atom representations per molecule.
• When implementing custom readout functions, always use the batch index rather than looping over molecules.
• For memory efficiency, set DataLoader with follow_batch if using nested batching (e.g., for drug-protein graphs).
• Large molecules (more than 100 atoms, macrocycles) or graphs with high connectivity can cause memory issues. Set a maximum number of atoms filter during preprocessing.

Further Reading

• Gilmer et al. (2017): MPNN for Quantum Chemistry
• Xiong et al. (2020): AttentiveFP
• PyTorch Geometric Documentation
• DGL Documentation
• Uni-Mol: Universal Molecular Representation
• Schutt et al. (2017): SchNet

Try it live

Explore a related model hands-on: open Graph Neural Networks: Learning on Relational Data and change its parameters while it runs. Everything stays in your browser.