A 2D molecular diffusion model built to understand how RFDiffusion and structure-based drug design work from first principles.
Modern protein and molecule design uses diffusion models — the same class behind image generators like Stable Diffusion — but operating over 3D atomic coordinates instead of pixels. RFDiffusion (Baker lab, 2023) generates novel protein backbones by learning to denoise backbone frames in SE(3) space.
This project builds the same ideas in 2D with small drug-like molecules, where everything is inspectable and trainable on a laptop. The goal is concrete intuition for why Gaussian noise works as a corruption process, what message passing computes, how denoising connects to valid molecular geometry, and how equivariance fixes the orientation-averaging failure mode that naive GNNs fall into.
See DESIGN.md for full technical details on the architecture, training, and design decisions.
~10k approved and clinical-stage small molecules from ChEMBL (max_phase ≥ 1), filtered to 4–50 heavy atoms, single connected components, valid RDKit SMILES. 2D coordinates generated by RDKit's chemistry-aware layout algorithm and normalized to mean bond length = 1.0.
An E(2)-equivariant GNN (EGNN) that takes noisy atom coordinates + bond graph and predicts clean coordinates (x0 prediction). Equivariance means the network reasons about atoms relative to their neighbors — rotating the molecule rotates the output identically — which prevents the naive GNN failure of averaging across all valid orientations into a blob at the centroid.
Diffusion runs over internal coordinates (bond lengths + absolute bond angles per spanning-tree edge) rather than raw Cartesian coordinates, so the noise directly corrupts the quantities that chemistry constraints are defined over.
After ~32 epochs on an M3 MacBook:
| Metric | Generated | Ground truth |
|---|---|---|
| Bond length std | 0.119 | 0.035 |
| Bond angle mean | 118° | 123° |
| Bond angle std | 41° | 24° |
The pre-EGNN baseline (80 epochs, non-equivariant GNN) had bond length std 0.525 and angle mean 65° — EGNN reached better geometry in fewer than half the epochs.
fetch_chembl.py fetch molecules from ChEMBL API → chembl_mols.json
dataset.py SMILES → 2D coords → Molecule2D dataclass, dataset stats
geom.py spanning tree, internal coordinates, constraint losses
noise.py cosine schedule, forward/reverse diffusion steps
model.py EGNN denoiser (equivariant message passing)
train.py training loop — supports internal coords + constraint loss
sample.py reverse diffusion sampler + constraint distribution plots
chembl_mols.json fetched dataset (~10k molecules)
python fetch_chembl.py # fetch dataset (one-time)
python train.py --use_internal --lambda_constraint 0.5
python sample.pyResume from a checkpoint:
python train.py --epochs 30 --resume best_model.pt --use_internal --lambda_constraint 0.5