BoKDiff: Best-of-K Diffusion Alignment for Target-Specific 3D Molecule Generation

TL;DR

BoKDiff tackles data scarcity and ligand-protein misalignment in structure-based drug design by introducing Best-of-K alignment to guide diffusion-based ligand generation. It builds on DecompDiff and combines multi-objective rewards over QED, SA, and Vina docking, plus Best-of-K alignment for fine-tuning and Best-of-N sampling for inference. On CrossDocked2020, BoKDiff achieves state-of-the-art results with high QED and docking performance, including a BoK-driven Vina score improvement and a 26% generation-success rate, while BoN demonstrates complementary robustness for property optimization. Limitations include potential distribution drift and the need for more sophisticated ligand relocation and multi-level optimization; future work targets KL-regularized rewards, broader docking metrics, and more flexible diffusion priors to further bridge generation and practical drug design.

Abstract

Structure-based drug design (SBDD) leverages the 3D structure of biomolecular targets to guide the creation of new therapeutic agents. Recent advances in generative models, including diffusion models and geometric deep learning, have demonstrated promise in optimizing ligand generation. However, the scarcity of high-quality protein-ligand complex data and the inherent challenges in aligning generated ligands with target proteins limit the effectiveness of these methods. We propose BoKDiff, a novel framework that enhances ligand generation by combining multi-objective optimization and Best-of-K alignment methodologies. Built upon the DecompDiff model, BoKDiff generates diverse candidates and ranks them using a weighted evaluation of molecular properties such as QED, SA, and docking scores. To address alignment challenges, we introduce a method that relocates the center of mass of generated ligands to their docking poses, enabling accurate sub-component extraction. Additionally, we integrate a Best-of-N (BoN) sampling approach, which selects the optimal ligand from multiple generated candidates without requiring fine-tuning. BoN achieves exceptional results, with QED values exceeding 0.6, SA scores above 0.75, and a success rate surpassing 35%, demonstrating its efficiency and practicality. BoKDiff achieves state-of-the-art results on the CrossDocked2020 dataset, including a -8.58 average Vina docking score and a 26% success rate in molecule generation. This study is the first to apply Best-of-K alignment and Best-of-N sampling to SBDD, highlighting their potential to bridge generative modeling with practical drug discovery requirements. The code is provided at https://github.com/khodabandeh-ali/BoKDiff.git.

Figures (2)

• Figure 1: The BoKDiff framework: From left to right, (1) Inputs: Randomly select a batch of the desired size from the training set. (2) Data Collection: Generate $K$ samples for each input pair. (3) Data Ranking: Compute the desired metrics—approximations of QED, SA, and the Vina Docking score—for the generated samples. Rank them using a weighted sum approach; this example emphasizes QED and SA for the final ranking. (4) Data Preparation: For the highest-reward sample, retrieve its docking pose and adjust its position by aligning its center of mass (CoM) to the CoM of the docking pose (See \ref{['preparation']} section). Perform final preprocessing by extracting the corresponding sub-pockets and sub-structures for the high-reward pairs. (5) Fine Tuning: Update the model parameters using a small learning rate and limited epochs. This process is demonstrated for the target protein 5ws0_B_rec.
• Figure 2: Success rates of reference and aligned models using the best-of-N strategy under various weight combinations for reward value determination. For all models, $N = 20$.