[2507.00445] Iterative Distillation for Reward-Guided Fine-Tuning of Diffusion Models in Biomolecular Design

Computer Science > Machine Learning arXiv:2507.00445 (cs) [Submitted on 1 Jul 2025 (v1), last revised 28 Feb 2026 (this version, v3)] Title:Iterative Distillation for Reward-Guided Fine-Tuning of Diffusion Models in Biomolecular Design Authors:Xingyu Su, Xiner Li, Masatoshi Uehara, Sunwoo Kim, Yulai Zhao, Gabriele Scalia, Ehsan Hajiramezanali, Tommaso Biancalani, Degui Zhi, Shuiwang Ji View a PDF of the paper titled Iterative Distillation for Reward-Guided Fine-Tuning of Diffusion Models in Biomolecular Design, by Xingyu Su and 9 other authors View PDF HTML (experimental) Abstract: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...