Reward Design for Post-Training Molecular Generators

Diffusion and flow-matching models can generate 3D molecules conditioned on protein pockets, but their training objectives do not directly optimize molecular properties relevant to practitioners. Post-training can shift the generative distribution toward molecules with better properties, yet its success depends on how rewards are constructed: molecular properties may be undefined for chemically invalid outputs, and optimizing one objective can degrade geometric plausibility or other desirable properties. We investigate these challenges by adapting DiffusionNFT, a recently proposed method for negative-aware finetuning, to jointly align the continuous coordinates and categorical outputs of a pocket-conditioned molecular generator. We first show that excluding chemically invalid generated samples from the alignment objective can cause validity collapse, whereas incorporating a validity reward largely prevents this failure. Building on this validity-aware formulation, we present the FlowR.NFT model family to systematically investigate combinations of chemical validity, Vina, QED, SA, and PoseBusters geometry rewards and their respective trade-offs. We identify reward compositions that improve Vina affinity, synthetic accessibility, and drug-likeness over the baseline model, with FlowR.NFT${}_{\mathrm{multi}}$ achieving a redocked Vina score of $-10.14kcal/mol$ on SPINDR while largely preserving molecular validity. These property gains transfer well to testing on CrossDocked2020, with a Vina docking score of $-9.10kcal/mol$, better than those of all prior baselines in our benchmark. However, the plausibility of generated ligand poses is less well preserved across datasets. Together, our results highlight the importance of careful invalid-sample handling and balancing molecular properties during fine-tuning of ligand generators.