MolJO

A medicinal chemist holding a hit compound bound in a crystal structure rarely wants a brand-new molecule. They want this molecule, with better affinity, without losing synthetic accessibility. Gradient guidance is the natural tool, but a 3D ligand is half continuous and half discrete: its coordinates take gradients cleanly, its atom types are categorical and do not. Guidance applied to coordinates alone pushes atoms into better contacts while never proposing that a carbon should have been a nitrogen, which is why coordinate-only methods improve docking scores while drug-likeness and synthesizability drift the wrong way.

MolJO (Molecule Joint Optimization) sidesteps this by guiding the model's belief rather than its samples. A Bayesian flow network carries a distribution over the clean molecule — a Gaussian over coordinates and a categorical simplex over atom types — both parameterized continuously and differentiably. One gradient of a property energy therefore reaches both modalities at once: it shifts the coordinate mean and reweights the type simplex in a single SE(3)-equivariant step, keeping the spatial and chemical halves of an edit consistent. A backward correction strategy then re-derives the belief from a sliding window of past steps instead of only the previous one, so successive gradients stay aligned and the window size trades exploration against exploitation.

The work comes from the Institute for AI Industry Research at Tsinghua University, with the Shanghai Institute of Materia Medica and Peking University, posted to arXiv in November 2024 and published at ICML 2025. Both halves run frozen at inference — a pocket-conditioned BFN backbone supplies the generative process, pretrained property predictors the gradients — so a new target needs no training. It sits alongside other Bayesian-flow treatments of pocket-conditioned design such as FusedBFN and MolPIF, which sample ligands de novo rather than improve an existing one.

#Key Features

• Joint multi-modal guidance: One property gradient updates continuous coordinates and discrete atom types together, so an affinity edit can change what an atom is, not only where it sits.
• Backward correction: Sampling optimizes over a sliding window of past beliefs, correcting earlier steps against the current best estimate rather than compounding a suboptimal history.
• Equivariance preserved: Guidance leaves sampling SE(3)-equivariant when the energy function is, so optimized poses do not depend on how the complex is oriented.
• Constrained editing: The same framework does R-group optimization, fragment growing, and scaffold hopping by infilling around a retained substructure.
• No oracle in the loop: Differentiable energies replace the docking simulations that resampling-based optimizers call on every round.

#Technical Details

MolJO uses CrossDocked2020, filtered to poses under 1 Å RMSD and clustered at 30% sequence identity, giving 100,000 training complexes and 100 test proteins; the property predictors train on the affinity-labelled version of that split, with released checkpoints for Vina Score and synthetic accessibility. Sampling runs 200 steps with a backward correction window of 130 and equal gradient weights on positions and types. Across the 100 test pockets it reports a 51.3% Success Rate (Vina Dock < -8.18, QED > 0.25, SA > 0.59), Vina Dock -9.05 and SA 0.78, against 11.1% and 0.69 for the coordinate-only baseline TAGMol; 39.8% of samples improve on the reference across all properties, versus under 17% for every 3D baseline. Ablations isolate both ingredients: backward correction moves guided sampling from Vina Dock -7.49 and SA 0.62 to -9.11 and 0.77, and guiding both modalities beats guiding either alone. Under structural constraints MolJO reaches a 29.0% success rate on R-group redesign while keeping 95.1% of generated molecules validly connected. Rescoring its top 10 samples with a docking oracle raises Success Rate to 70.3%.

#Applications

The practical target is lead optimization against a solved structure: supply a protein PDB file and a reference ligand, and MolJO returns analogues optimized jointly for binding affinity, drug-likeness and synthesizability. Because the retained substructure is a constraint, it covers the edits medicinal chemistry actually asks for — growing an R-group into an adjacent subpocket, swapping a scaffold while preserving a key interaction, or trading affinity against synthetic tractability by reweighting objectives. Cheminformatics groups can also attach a differentiable property model of their own to the frozen backbone.

#Impact

MolJO makes gradient guidance a workable alternative to oracle-driven resampling for structure-based optimization, and its joint continuous-discrete formulation is the piece that generalizes: the argument applies to any generative model whose intermediate state is a differentiable belief rather than a noisy sample. Evaluation is entirely in silico, with Vina docking scores, QED and SA as proxies for the properties that matter, and no experimental validation of the optimized compounds. Code and checkpoints carry a CC BY-NC-SA license, ruling out commercial use, and sit in a shared repository alongside two sibling projects from the same group, so the MolJO subdirectory and its own checkpoints are what apply here.