BoKDiff

Diffusion models for structure-based drug design learn to place atoms inside a protein binding pocket, but they are trained on a likelihood objective that has nothing to say about whether the resulting molecule is drug-like, synthesizable, or actually binds. The usual result is a generator that reproduces the statistics of its training complexes rather than optimizing the properties a medicinal chemist cares about — and the supply of high-quality protein-ligand structures is too small to fix this by scaling data.

BoKDiff, from Ali Khodabandeh Yalabadi, Mehdi Yazdani-Jahromi, and Ozlem Ozmen Garibay at the University of Central Florida, imports the Best-of-K alignment recipe used to post-train large language and vision models into this setting. It starts from DecompDiff, a pocket-conditioned 3D diffusion generator, samples many candidates per pocket, scores each on a weighted combination of QED, synthetic accessibility, and Vina docking, keeps the top-ranked candidate per pocket, and fine-tunes the diffusion model on that filtered set. Iterating this loop yields a distinct BoKDiff checkpoint, trained once on the CrossDocked2020 training pockets and then applied unchanged to held-out test pockets.

A second, lighter contribution is Best-of-N sampling: the same ranking function applied purely at inference time, selecting one molecule from N candidates without touching the weights. The paper appeared as a preprint in January 2025 and was published in Bioinformatics Advances later that year.

#Key Features

• Best-of-K alignment fine-tuning: For each pocket the model generates K candidates, ranks them by a weighted reward over QED, synthetic accessibility, and docking score, and fine-tunes on the winners at a learning rate of 1e-6 for 1,000 iterations per round — a rejection-sampling analogue of preference alignment, applied to a molecular diffusion model.
• Centre-of-mass realignment: Generated ligands are repositioned so their centre of mass matches that of their docking pose before decomposition, which is what makes the extracted arms and scaffold correspond to the pose the reward was computed on.
• Multi-objective reward weighting: The QED / SA / docking weights are exposed as a configuration parameter, so a run can be biased toward drug-likeness or toward affinity.
• Training-free Best-of-N mode: The same ranking applied over N=30 samples at inference, with no weight updates, reaches a 37.76% success rate with mean QED 0.67 and SA 0.76.

#Technical Details

The backbone is DecompDiff's decomposed-prior diffusion model: an equivariant graph transformer denoiser operating on a k-nearest-neighbour graph (k=32) over pocket and ligand atoms, with 6 layers, hidden dimension 128, 16 attention heads, 1,000 diffusion timesteps, and joint diffusion over atom positions, atom types, and 5 bond classes. Alignment data collection sampled 10 molecules for each of roughly 100,000 complexes, and fine-tuning ran once per alignment iteration rather than per sample.

On the CrossDocked2020 test split, the released BoKDiff checkpoint reports a mean Vina Dock of -8.58, Vina Min of -7.50, Vina Score of -5.92, QED 0.48, SA 0.60, and a 26.0% success rate, against DecompDiff's -8.39 / -7.04 / -5.67 and 24.5%, and reference ligands at -7.45 and 25.0%. The abstract's headline "-8.58 average docking score" is the Vina Dock column of Table 1, the redocked score; the direct Vina Score of the generated pose is -5.92. The Best-of-N variant trades affinity for chemistry, reaching Vina Dock -8.06 with QED 0.67 and a 37.76% success rate. The authors note that DecompDiff's fixed atom-count constraint bounds the complexity of what can be generated, that centre-of-mass repositioning is a crude alignment, and that optimizing a single docking program risks overfitting to its scoring function.

#Applications

BoKDiff addresses hit generation in structure-based drug discovery: given a resolved or modelled binding site, it proposes complete 3D ligand geometries whose ranking already accounts for synthesizability and predicted affinity, reducing the filtering burden before candidates go to docking triage or synthesis planning. The Best-of-N mode is the more immediately portable idea, since it can be wrapped around an existing pocket-conditioned generator without retraining. Code and both checkpoints — the base DecompDiff weights and the aligned BoKDiff model — are distributed through the repository and linked Google Drive folders under a CC BY-NC 4.0 licence, which permits research use but not commercial deployment.

#Impact

The contribution is methodological rather than architectural: BoKDiff shows that the alignment machinery developed for language models transfers to molecular diffusion, where the reward is a cheap in-silico property estimate rather than a human preference label. That matters because scoring a generated molecule is far cheaper than obtaining another crystallographic complex, so reward-guided fine-tuning is a way to keep improving after the structural data runs out. The gains over the DecompDiff baseline are real but modest, and all evaluation is in-silico on a single benchmark with Vina as both the reward and the metric — no experimental validation of any generated compound is reported.