Generalizing Equivariant Neural Networks for Protein–Ligand Binding by Comparing Atomic Environments
Abstract
Structure-based models for protein--ligand binding affinity often owe their accuracy to close analogs shared between training and test sets, and generalize poorly when those analogs are removed. We study which architectural choices let an equivariant message-passing network generalize under decontaminated evaluation, using only the experimental affinities of PDBbind CleanSplit and no pretrained inputs. Our model, MPNN+PR (a message-passing network with paired readouts), processes a bound and an edge-masked unbound copy of each complex and compares their message-passing outputs at the atom, edge, and pocket levels through operations we term \emph{paired readouts}. The atom-graph model is the most accurate method on CASF-2016 (RMSE1.258) after training on CleanSplit. Because binding affinity is an ensemble quantity, we hypothesize that a network with higher angular-order features can improve by also being robust to variations in the atomic positions. We pursue this through two changes: computing higher-order features from a separate spherical-harmonic expansion of the local atomic environment, and coarse-graining the pocket into residues. Both reduce the model's sensitivity to coordinate noise. A residue variant has the highest correlation on the independent subset, and after finetuning on a few labeled complexes it outperforms prior models on held-out pocket clusters.