Ensemble-Conditioned Molecular Design
Ross Irwin, Alessandro Tibo, Jon Paul Janet, Simon Olsson
cs.LG, cs.NE
2026-09-14
A 3D flow-matching generator trains on single conditions, then composes vector fields at inference to target or avoid ensemble modes. Dual AChE/MAO-B binders rise from 9% to 23%.
Computational molecular design still scores a molecule by one pose. Docking, shape-based virtual screening, and both structure-based and ligand-based generators train and evaluate on the bioactive conformer. Real molecules live on a Boltzmann distribution. Oral bioavailability, membrane permeability, the entropic cost of binding, and whether a ligand is an agonist or an antagonist all depend on that distribution.
Dual-target design needs a molecule that can occupy bioactive poses in two pockets. Agonist design needs it to stabilise the active receptor state and stay off the inactive one. Selectivity is the inverse. Current generators take one condition at a time. Search-based methods re-run for every new combination, and they get worse as orthogonal terms pile up. This paper recasts design as joint optimisation of ensemble modes and ensemble properties, and moves the composition to inference.
The method is ensemble-conditioned guidance on a 3D flow-matching generator. Coordinates use continuous flow matching; atom and bond types use discrete flow matching. The network predicts endpoints.
Modes are encoded as noisy coordinate clouds for shape, pharmacophore point clouds, or protein pockets taken as residues within 6 Å of a reference ligand. Properties are scalars over a precomputed ensemble: mean polar surface area (PSA) and mean pairwise RMSD. Training uses only single conditions, drawn from GEOM Drugs (300k molecules with CREST ensembles) and SPINDR (35k protein-ligand complexes). Conditions are independently masked, including a joint drop that supplies the null input for classifier-free guidance.
At sampling, each mode condition gets its own forward pass. Endpoint predictions are mixed with weights αk: linear mix on coordinates, log-linear on categoricals. Positive αk steers toward that mode; negative αk steers away. A global scale γ plays the usual CFG role.
Two supporting tricks make composition usable. Adaptive symmetry learning randomly rotates encoder inputs so the same encoder can emit E(3)-equivariant or invariant features. At inference, one condition defines the generation frame (equivariant); the rest are encoded invariantly, which stops equivariant vectors from different frames from twisting the pose. Flexible size is handled by padding to 48 heavy atoms, placing pad atoms at the centre of mass, and permuting prior atoms to match. Unconditional samples average 22.1 heavy atoms against a training mean of 24.8. The work comes from AstraZeneca Molecular AI and Chalmers. Training ran 200 epochs on one A100, about two days. Sampling uses 100 Euler steps.
They built two benchmarks because almost none existed.
On multi-mode shape, compact and extended conformers are paired from the GEOM Drugs test set: 496 pairs for targeting both, 535 for targeting one and avoiding the other. With γ=4.0 and σshape=0.2, the equivariant side gains about 0.19 shape Tanimoto over a size-matched virtual screening baseline; the invariant side gains 0.07. Connected validity stays above 0.97. Raising σshape to 0.5 drops mean ECFP Tanimoto to the reference molecules from 0.25-0.31 to 0.14. Fidelity versus novelty is a knob.
On property control, 179 held-out SPINDR systems are conditioned on pocket plus crystal pharmacophores, then given a PSA or RMSD target.
| Condition | PSA (Ų) | RMSD (Å) | Interaction recovery | Vina (kcal/mol) |
| Pocket + pharma | 167 | 2.07 | 0.954 | -6.18 |
| PSA target 80 to 140 | 111 to 163 | 2.11-2.18 | 0.944-0.958 | -6.17 to -6.31 |
| RMSD target 1.0 to 2.5 | 154-159 | 1.69 to 2.17 | 0.954-0.961 | -6.90 to -6.18 |
Control is directional. Values compress toward the unconditioned baseline at both ends of each sweep.
Two drug-design case studies:
| Task | Setup | Outcome |
| Dual AChE / MAO-B | Pharmacophores of donepezil and safinamide together | 23% dual binders within 1.0 kcal/mol of both drugs, vs 9% and 7% for single-target controls |
| A2A agonist | Active pocket positive, inactive pocket negative | Mean ΔVinastate from +0.47 (active only) to +0.97 / +1.26 kcal/mol |
The best dual candidate scores -12.24 on AChE (donepezil -12.19) and -11.85 on MAO-B (safinamide -10.31), overlaying each drug's ring system on a new scaffold. At the 0.5 kcal/mol cutoff the dual rate is 10%, against 1% and 3%.
For people building 3D generators, multi-state design stops being a data problem and becomes a sampling combination. No paired multi-mode labels are required. Shapes, pharmacophores, pockets, and ensemble scalars can be added or subtracted at inference. Negative weights restore an "avoid this pocket" interface that most recent multi-state protein methods dropped.
Single-pocket, single-shape generation is already crowded. The missing piece is saying "want this, not that, and push PSA down" in one sample. The costs are explicit: runtime grows linearly with the number of modes, the invariant side is weaker, and property control has direction without calibration. Experimental success is still untested. The oracle is docking.
The authors list them. Invariant conditions carry a weaker signal than the equivariant frame-defining one, so which condition owns the frame changes the result. Inference cost scales linearly with mode count. Property conditioning is directional, compressed toward the baseline at the sweep ends. Evaluation uses ETKDG plus MMFF94, not CREST, so PSA and RMSD only approximate the training observables. Case studies use AutoDock Vina as a binding oracle. Heavy-atom cap is 48, which rules out chameleonic PROTACs and large macrocycles. The multi-mode benchmark is essentially two modes.
Docking as the verdict on dual-target and agonist design is the number that needs the biggest discount. 23% and +1.26 kcal/mol live in Vina's world.