A Graph Completion Method that Jointly Predicts Geometry and Topology Enables Effective Molecule Assembly

Abstract

A common starting point for drug design is to find small chemical groups or "fragments" that form interactions with distinct subregions in a protein binding pocket. The subsequent challenge is to assemble these fragments into a molecule that has high affinity to the protein, by adding chemical bonds between atoms in different fragments. This "molecule assembly" task is particularly challenging because, initially, fragment positions are known only approximately. Prior methods for spatial graph completion-adding missing edges to a graph whose nodes have associated spatial coordinates-either treat node positions as fixed or adjust node positions before predicting edges. The fact that these methods treat geometry and topology prediction separately limits their ability to reconcile noisy geometries and plausible connectivities. To address this limitation, we introduce EdGr, a spatial graph diffusion model that reasons jointly over geometry and topology of molecules to simultaneously predict fragment positions and inter-fragment bonds. Importantly, predicted edge likelihoods directly influence node position updates during the diffusion denoising process, allowing connectivity cues to guide spatial movements, and vice versa. EdGr substantially outperforms previous methods on the molecule assembly task and maintains robust performance as noise levels increase. Beyond drug discovery, our approach of explicitly coupling geometry and topology prediction is broadly applicable to spatial graph completion problems, such as neural circuit reconstruction, 3D scene understanding, and sensor network design.