# Fine-Tuning Diffusion Models for Molecular Generation via Reinforcement Learning and Fast Sampling

arXiv: https://arxiv.org/abs/2606.01220 · HTML: https://mlsanity.com/paper/2606.01220/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Flow matching, time series and science, one of 14 directions in Diffusion models.
- Builds on 17 of its 30 references that are in the map.
- First posted on arXiv 31 May 2026.

## Placed in

- Diffusion models / [Flow matching, time series and science](https://mlsanity.com/diffusion/flow-matching-time-series-and-science/)
- Deep reinforcement learning methods / [RL for energy, networks and operations](https://mlsanity.com/deeprl/rl-for-energy-networks-and-operations/)

## Builds on (most cited first, 17 of 17)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [UMAP: Uniform Manifold Approximation and Projection for Dimension Reduction](https://mlsanity.com/paper/1802.03426/) (2018)
- [Training Diffusion Models with Reinforcement Learning](https://mlsanity.com/paper/2305.13301/) (2023)
- [Diffusion Model Alignment Using Direct Preference Optimization](https://mlsanity.com/paper/2311.12908/) (2023)
- [Structure-based Drug Design with Equivariant Diffusion Models](https://mlsanity.com/paper/2210.13695/) (2022)
- [Pocket2Mol: Efficient Molecular Sampling Based on 3D Protein Pockets](https://mlsanity.com/paper/2205.07249/) (2022)
- [3D Equivariant Diffusion for Target-Aware Molecule Generation and Affinity Prediction](https://mlsanity.com/paper/2303.03543/) (2023)
- [A 3D Generative Model for Structure-Based Drug Design](https://mlsanity.com/paper/2203.10446/) (2022)
- [Generating 3D Molecules for Target Protein Binding](https://mlsanity.com/paper/2204.09410/) (2022)
- [DecompDiff: Diffusion Models with Decomposed Priors for Structure-Based Drug Design](https://mlsanity.com/paper/2403.07902/) (2024)
- [MolCRAFT: Structure-Based Drug Design in Continuous Parameter Space](https://mlsanity.com/paper/2404.12141/) (2024)
- [DecompOpt: Controllable and Decomposed Diffusion Models for Structure-based Molecular Optimization](https://mlsanity.com/paper/2403.13829/) (2024)
- [Aligning Target-Aware Molecule Diffusion Models with Exact Energy Optimization](https://mlsanity.com/paper/2407.01648/) (2024)
- [Target-aware Molecular Graph Generation](https://mlsanity.com/paper/2202.04829/) (2022)
- [TAGMol: Target-Aware Gradient-guided Molecule Generation](https://mlsanity.com/paper/2406.01650/) (2024)

## Abstract (arXiv)

Generating molecules that simultaneously satisfy drug-like properties and conform to the 3D structure of a target protein is a core challenge in structure-based drug design (SBDD). Existing generative approaches, however, often rely on costly post-hoc processing during Sampling or require carefully curated datasets during training, yet still achieve modest gains. These limitations are especially pronounced in multi-objective settings, where balancing conflicting criteria remains a core challenge. To address these challenges, We propose FTDiff, a reinforcement learning fine-tuning framework tailored for diffusion-based molecular generation under structural constraints. To ensure stable and sample-efficient optimization, FTDiff adopts a group relative policy optimization (GRPO) style strategy. Furthermore, FTDiff builds upon a time-free pretrained diffusion model and incorporates a fast sampling mechanism that reduces the number of denoising steps, significantly accelerating both training and inference while maintaining generation quality. By optimizing a fixed threshold-aware reward, FTDiff effectively guides the model to produce valid, diverse, and high- quality molecules that balance multiple drug design objectives. Extensive experiments on benchmark datasets demonstrate that FTDiff consistently outperforms prior methods, without requiring expensive post-hoc optimization or intricate data engineering.