# Iterative Distillation for Reward-Guided Fine-Tuning of Diffusion Models in Biomolecular Design

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

## Key facts

- Cited by 11 papers in the map (data through 14 Oct 2026).
- Placed in Flow matching, time series and science, one of 14 directions in Diffusion models.
- Builds on 43 of its 71 references that are in the map.
- First posted on arXiv 1 Jul 2025.

## Placed in

- Diffusion models / [Flow matching, time series and science](https://mlsanity.com/diffusion/flow-matching-time-series-and-science/)

## Builds on (most cited first, 30 of 43)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Training Diffusion Models with Reinforcement Learning](https://mlsanity.com/paper/2305.13301/) (2023)
- [Simple and Effective Masked Diffusion Language Models](https://mlsanity.com/paper/2406.07524/) (2024)
- [Advantage-Weighted Regression: Simple and Scalable Off-Policy Reinforcement Learning](https://mlsanity.com/paper/1910.00177/) (2019)
- [Improving Diffusion Models for Inverse Problems using Manifold Constraints](https://mlsanity.com/paper/2206.00941/) (2022)
- [DPOK: Reinforcement Learning for Fine-tuning Text-to-Image Diffusion Models](https://mlsanity.com/paper/2305.16381/) (2023)
- [Directly Fine-Tuning Diffusion Models on Differentiable Rewards](https://mlsanity.com/paper/2309.17400/) (2023)
- [Universal Guidance for Diffusion Models](https://mlsanity.com/paper/2302.07121/) (2023)
- [Aligning Text-to-Image Diffusion Models with Reward Backpropagation](https://mlsanity.com/paper/2310.03739/) (2023)
- [Score-based Generative Modeling of Graphs via the System of Stochastic Differential Equations](https://mlsanity.com/paper/2202.02514/) (2022)
- [Inference-Time Scaling for Diffusion Models beyond Scaling Denoising Steps](https://mlsanity.com/paper/2501.09732/) (2025)
- [Practical and Asymptotically Exact Conditional Sampling in Diffusion Models](https://mlsanity.com/paper/2306.17775/) (2023)
- [A General Framework for Inference-time Scaling and Steering of Diffusion Models](https://mlsanity.com/paper/2501.06848/) (2025)
- [Beyond Reverse KL: Generalizing Direct Preference Optimization with Diverse Divergence Constraints](https://mlsanity.com/paper/2309.16240/) (2023)
- [Derivative-Free Guidance in Continuous and Discrete Diffusion Models with Soft Value-Based Decoding](https://mlsanity.com/paper/2408.08252/) (2024)
- [Unlocking Guidance for Discrete State-Space Diffusion and Flow Models](https://mlsanity.com/paper/2406.01572/) (2024)
- [Aligning Language Models with Preferences through f-divergence Minimization](https://mlsanity.com/paper/2302.08215/) (2023)
- [Exploring Chemical Space with Score-based Out-of-distribution Generation](https://mlsanity.com/paper/2206.07632/) (2022)
- [Test-time Alignment of Diffusion Models without Reward Over-optimization](https://mlsanity.com/paper/2501.05803/) (2025)
- [Inference-Time Alignment in Diffusion Models with Reward-Guided Generation: Tutorial and Review](https://mlsanity.com/paper/2501.09685/) (2025)
- [Distilling Policy Distillation](https://mlsanity.com/paper/1902.02186/) (2019)
- [Fine-Tuning Discrete Diffusion Models via Reward Optimization with Applications to DNA and Protein Design](https://mlsanity.com/paper/2410.13643/) (2024)
- [Large-scale Reinforcement Learning for Diffusion Models](https://mlsanity.com/paper/2401.12244/) (2024)
- [A general sample complexity analysis of vanilla policy gradient](https://mlsanity.com/paper/2107.11433/) (2021)
- [Amortizing intractable inference in diffusion models for vision, language, and control](https://mlsanity.com/paper/2405.20971/) (2024)

## Built on by (most cited first, 11 of 11)

- [d2: Improving Reasoning in Diffusion Language Models via Trajectory Likelihood Estimation](https://mlsanity.com/paper/2509.21474/) (2025)
- [TR2-D2: Tree Search Guided Trajectory-Aware Fine-Tuning for Discrete Diffusion](https://mlsanity.com/paper/2509.25171/) (2025)
- [Autonomous Agents for Scientific Discovery: Orchestrating Scientists, Language, Code, and Physics](https://mlsanity.com/paper/2510.09901/) (2025)
- [Diffusion Alignment as Variational Expectation-Maximization](https://mlsanity.com/paper/2510.00502/) (2025)
- [Generative Modeling in Protein Design: Neural Representations, Conditional Generation, and Evaluation Standards](https://mlsanity.com/paper/2603.26378/) (2026)
- [Seek and You Shall Fold](https://mlsanity.com/paper/2511.13244/) (2025)
- [On the Design Space of Discrete Diffusion Online Adaptation for Molecular Optimization](https://mlsanity.com/paper/2607.02834/) (2026)
- [Unifying ICL, SFT, KL-Regularized RL Through a Bayesian Lens](https://mlsanity.com/paper/2609.05111/) (2026)
- [Graph Forward Distribution Matching for Molecular Inverse Design](https://mlsanity.com/paper/2609.32056/) (2026)
- [FestDPO: Few-step Generator Alignment with Direct Preference Optimization](https://mlsanity.com/paper/2609.34673/) (2026)
- [LLM-to-Flow Distillation: Teaching Natural Language Goals to Scientific Generators](https://mlsanity.com/paper/2610.17246/) (2026)

## Abstract (arXiv)

We address the problem of fine-tuning diffusion models for reward-guided generation in biomolecular design. While diffusion models have proven highly effective in modeling complex, high-dimensional data distributions, real-world applications often demand more than high-fidelity generation, requiring optimization with respect to potentially non-differentiable reward functions such as physics-based simulation or rewards based on scientific knowledge. Although RL methods have been explored to fine-tune diffusion models for such objectives, they often suffer from instability, low sample efficiency, and mode collapse due to their on-policy nature. In this work, we propose an iterative distillation-based fine-tuning framework that enables diffusion models to optimize for arbitrary reward functions. Our method casts the problem as policy distillation: it collects off-policy data during the roll-in phase, simulates reward-based soft-optimal policies during roll-out, and updates the model by minimizing the KL divergence between the simulated soft-optimal policy and the current model policy. Our off-policy formulation, combined with KL divergence minimization, enhances training stability and sample efficiency compared to existing RL-based methods. Empirical results demonstrate the effectiveness and superior reward optimization of our approach across diverse tasks in protein, small molecule, and regulatory DNA design. The source code is released at (https://divelab.github.io/VIDD/).