Equ-GNN

E(3)-equivariant graph neural networks for protein–ligand binding-affinity prediction.

Two equivariant architectures — an EGNN baseline and an e3nn tensor-product model — are built on a shared, reproducible PyG data pipeline, verified for E(3) invariance to numerical precision, and trained/evaluated on a small PDBbind-style subset. This repository is a correct, documented reference implementation, sized to run end-to-end on a CPU; the path to competitive accuracy on full PDBbind (GPU) is described in Scaling to GPU.

---

1. Task and the invariance contract

Task. Regress the scalar binding affinity pKd = -log10(K_d / K_i / IC50 in molar) of a protein–ligand complex from its 3D structure. Higher pKd = tighter binding.

Why this task. A scalar target gives the cleanest possible statement of the equivariance principle central to this field: the physical quantity does not depend on the orientation or position of the complex in space, so the prediction must be invariant, while the internal directional features the network builds must be equivariant. This is the same distinction as energy is invariant, force is equivariant in ML interatomic potentials (NequIP, MACE).

The contract (enforced in tests/test_equivariance.py):

Quantity	Must transform as	Verified deviation (float64)
Scalar output pKd, under rotation R ∈ SO(3) + translation t	Invariant: f(Rx+t) = f(x)	EGNN 1.8e-14, e3nn 2.6e-15
Scalar output, under reflection (O(3))	Invariant	0.0 (both)
e3nn internal l=1 (1o) features, under R	Equivariant: h(Rx) = D(R) h(x)	1.7e-8

The l=1 equivariance floor (~1e-7 at float64) is set by e3nn.o3.spherical_harmonics's polynomial evaluation, verified in isolation — it is a library precision limit, not a model defect. It does not affect the prediction, because the scalar readout consumes only the 0e channels (which match to ~1e-13). Binding affinity is treated as parity-invariant; a chirality-sensitive target would require breaking parity (odd vs even irreps, 1e vs 1o).

How the guarantee is achieved — architecturally, not learned. Atomic coordinates never enter the network as raw features. They are stored separately in data.pos and enter only through relative vectors x_i − x_j and their invariants (squared distance for EGNN; spherical harmonics of the direction + a radial basis of the length for e3nn). Because every layer is an equivariant map, invariance/equivariance holds for any weights, before and after training.

---

2. Repository layout

equ_gnn/
├── configs/default.yaml       # single source of truth for every run (seed, data, model, training)
├── src/
│   ├── data_fetch.py          # build labeled manifest from RCSB (structures + Kd/Ki/IC50)
│   ├── graphs.py              # complex -> pocket-cropped atom graph (PyG Data)
│   ├── dataset.py             # fetch all structures, cache PyG graphs
│   ├── model_egnn.py          # EGNN baseline (Satorras et al. 2021)
│   ├── model_e3nn.py          # e3nn tensor-product model (NequIP/MACE style)
│   ├── train.py               # seed-fixed training + overfit sanity check + metrics
│   └── evaluate.py            # predicted-vs-true plots
├── tests/test_equivariance.py # E(3) invariance / equivariance unit tests
├── data/                      # manifest.csv, cached graphs.pt, raw/*.cif
├── results/                   # equivariance_report.json, results.json, pred_vs_true.png
└── environment.txt            # pinned versions (CPU)

Run order (from equ_gnn/):

python src/data_fetch.py         # -> data/manifest.csv
python src/dataset.py            # -> data/graphs.pt
python tests/test_equivariance.py# -> results/equivariance_report.json  (PASS)
python src/train.py              # -> results/results.json, *_weights.pt
python src/evaluate.py           # -> results/pred_vs_true.png

---

3. Data pipeline (design decisions)

• Source: RCSB PDB REST API, not the PDBbind archive. PDBbind's official download is registration-gated. RCSB annotates the same affinity measurements (rcsb_binding_affinity, provenance BindingDB / PDBbind / Binding MOAD) on the identical structures and is openly fetchable, giving a fully reproducible pipeline with no manual download. The search API returns ~1,600 Kd-typed, ≤2.5 Å, protein-containing complexes; the current subset takes the first 30 with cleanly parseable single-ligand labels.
• Label: pKd = -log10(K in molar), unit-converted from the reported nM/µM/etc. The subset spans pKd 2.9–9.9 (≈7 log units) — a well-spread regression target.
• Pocket cropping (pocket_cutoff = 8 Å). A whole protein is thousands of atoms; affinity is a local, pocket-level property. Keeping only protein atoms within 8 Å of any ligand atom yields graphs of 123–488 atoms (mean 261) — small enough for CPU, retaining the informative geometry.
• Graph (edge_cutoff = 5 Å, max_neighbors = 32). A radius graph over heavy atoms (H removed, water removed). Nodes carry scalar features only: element one-hot + ligand/protein flag. A pure-torch radius graph avoids the fragile compiled torch-cluster dependency.

---

4. Models (design decisions)

Both map a Data object to one scalar and use permutation- and size-invariant mean pooling (pocket graphs vary widely in atom count; sum pooling would make the output scale track graph size).

EGNN baseline (model_egnn.py; Satorras, Hoogeboom & Welling, ICML 2021). Messages depend on node features and the squared distance ‖x_i − x_j‖² — an E(3) invariant — so scalar features stay invariant by construction, with no spherical harmonics or tensor products. The coordinate update is disabled: for an invariant predictor we only need the scalar readout, and freezing coordinates keeps the forward map a pure invariant function of geometry. LayerNorm on the scalar features and degree-normalized aggregation keep activations well-scaled. ~80k parameters.

e3nn tensor-product model (model_e3nn.py; NequIP — Batzner et al. 2022; MACE — Batatia et al. 2022; e3nn — Geiger & Smidt 2022). Genuinely directional: hidden features carry l=0 (scalar) and l=1 (vector) irreps (32x0e + 8x1o). Per edge, the relative direction becomes spherical harmonics (equivariant), the distance becomes a radial basis feeding an MLP that produces the tensor-product path weights (invariant), and a Clebsch–Gordan FullyConnectedTensorProduct mixes node features with the harmonics. An equivariant Gate nonlinearity between layers provides expressivity that a purely (bi)linear tensor-product stack lacks. ~180k parameters.

Why two models. EGNN uses only scalar invariants of the geometry (distances); the e3nn model carries true angular information (l≥1) through equivariant tensor products. Comparing them isolates what full E(3)-equivariant message passing buys over a distance-only baseline — the central scientific question of the project.

---

5. Training and results

Seed-fixed (42), Adam, MSE on z-scored labels (train statistics only; metrics reported in original pKd units), early stopping on validation RMSE, global grad-norm clipping (guards against rare tensor-product magnitude spikes). A single-batch overfit sanity check runs first: a model that cannot drive one batch's loss toward zero has a bug independent of generalization. Both pass (EGNN 0.26, e3nn 0.26, no divergence).

Held-out test metrics (subset of 300 → ~210 train / 45 val / 45 test):

Model	Epochs	RMSE (pKd)	Pearson r	Spearman ρ
EGNN (distance-only)	107 (early-stopped)	1.53	0.56	0.53
e3nn tensor-product	15 (CPU-budget-limited)	1.81	0.17	0.12

How to read this. These numbers come from a 300-complex subset — 10× the initial 30 — which is what turned noise into signal. At n=30 both models scored Pearson ≈ 0.2 on a 5-complex test set (meaningless); at n=300 the EGNN baseline reaches Pearson 0.56 / Spearman 0.53 on 45 held-out complexes from diverse protein families, a genuine correlation. This confirms the earlier flat results were a data-scale limitation, not a code defect.

The e3nn model is under-trained, not broken: its overfit-one-batch check passes (0.26, identical to EGNN), so it learns — but its tensor-product convolutions cost ~134 s/epoch on CPU vs ~5 s/epoch for EGNN (~25×), so within a practical CPU budget it ran only 15 epochs and still predicts near the training mean (flat line in results/pred_vs_true.png). This is the expected CPU trade-off: the more expressive equivariant model needs GPU-scale compute to converge. A like-for-like comparison (equal epochs) requires the GPU path below. Do not read the EGNN-vs-e3nn gap here as evidence that tensor products underperform — it reflects unequal training budgets, not model capacity.

Competitive accuracy (literature CASF-2016 Pearson ≈ 0.8) requires the full PDBbind (~19k complexes) on GPU. The deliverables that are fully validated here — reproducible pipeline, two correct equivariant models, machine-precision invariance, and a working (r = 0.56) baseline predictor — are complete.

---

6. Scaling to GPU

The code is device-agnostic (cfg.device); nothing here assumes CPU. To scale:

1. Data. Replace the 300-row manifest with the full PDBbind v2020 general+refined sets (~19k complexes), or expand the RCSB query (subset_size, drop the resolution filter). Use the CASF-2016 core set as a held-out test split for comparability with the literature.
2. Precompute + batch. Cache graphs to disk (already done via graphs.pt; for large N use an on-disk InMemoryDataset or LMDB). Increase batch_size and use num_workers > 0.
3. Device. Set device: cuda, move model and batch to GPU (.to(device)), enable AMP (torch.cuda.amp) for the e3nn model — tensor products are the compute bottleneck.
4. torch-cluster (optional). At scale, install the CUDA torch-cluster wheel matching the torch/CUDA version and swap the pure-torch radius graph in graphs.py for torch_geometric.nn.radius_graph.
5. Model capacity. Add irreps (lmax=2, higher multiplicities), more layers, and a learned radial cutoff. Consider MACE-style higher body-order messages.
6. Remote dispatch. On this platform, add a GPU host in the Compute panel and submit train.py there; keep parsing/plotting local.

The equivariance test should be re-run after any architectural change — it is the invariant that must never regress.

---

7. References

• Satorras, Hoogeboom, Welling. E(n) Equivariant Graph Neural Networks. ICML 2021.
• Batzner et al. E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials (NequIP). Nat. Commun. 2022.
• Batatia et al. MACE: Higher Order Equivariant Message Passing Neural Networks. NeurIPS 2022.
• Geiger & Smidt. e3nn: Euclidean Neural Networks. 2022. arXiv:2207.09453.
• Liu et al. PDBbind / Su et al. CASF-2016. J. Chem. Inf. Model. 2019.
• Berman et al. The Protein Data Bank. Nucleic Acids Res. 2000. (RCSB REST API data source.)