Empower Structure-Based Molecule Optimization with Gradient Guided Bayesian Flow Networks

Structure-based molecule optimization (SBMO) aims to optimize molecules with both continuous coordinates and discrete types against protein targets.A promising direction is to exert gradient guidance on generative models given its remarkable success in images, but it is challenging to guide discrete data and risks inconsistencies between modalities.To this end, we leverage a continuous and differentiable space derived through Bayesian inference, presenting Molecule Joint Optimization (MolJO), the gradient-based SBMO framework that facilitates joint guidance signals across different modalities while preserving SE(3)-equivariance.We introduce a novel backward correction strategy that optimizes within a sliding window of the past histories, allowing for a seamless trade-off between explore-and-exploit during optimization.MolJO achieves state-of-the-art performance on CrossDocked2020 benchmark (Success Rate 51.3\%, Vina Dock -9.05 and SA 0.78), more than 4x improvement in Success Rate compared to the gradient-based counterpart, and 2x ``Me-Better'' Ratio as much as 3D baselines. Furthermore, we extend MolJO to a wide range of optimization settings, including multi-objective optimization and challenging tasks in drug design such as R-group optimization and scaffold hopping, further underscoring its versatility.Code is available at https://github.com/AlgoMole/MolCRAFT.

Problem: Designing molecules for new drugs is challenging because we need to optimize both the molecule’s shape (where atoms are positioned, continuous) and their chemical components (what types of atoms to use, discrete). Traditional approaches struggle to effectively and smoothly guide both aspects together. Solution: We developed MolJO, a framework that creates a unified, flexible “workspace” where a molecule’s 3D structure and chemical parts can be optimized together. Think of it as a molecular sculptor that adjusts both the sculpture’s shape and material in real time. A novel “memory window” technique balances exploring new designs and refining existing ones.Impact: - 4× better success rate than previous gradient-based methods (51.3%). - Produces molecules with stronger binding (-9.05 docking score) and safer chemical profiles (SA score 0.78). - Adapts to practical tasks like R-group design and scaffold hopping.