迭代蒸馏：用于生物分子设计中扩散模型的奖励引导微调
机器学习
2026-03-03 v3 人工智能
定量方法
摘要
我们解决了扩散模型在生物分子设计中进行奖励引导生成的微调问题。尽管扩散模型在建模复杂高维数据分布方面效果卓越，但实际应用往往需要超越高保真生成，要求针对可能非微分的奖励函数进行优化，例如基于物理仿真或基于科学知识的奖励。虽然已探索 RL 方法对扩散模型进行微调以实现此目标，但常因其在策略上的特性导致训练不稳定、样本效率低和模式崩溃。本文提出一种基于迭代蒸馏的微调框架，使扩散模型能够针对任意奖励函数进行优化。我们将问题建模为策略蒸馏：在 roll-in 阶段收集离策略数据，在 roll-out 阶段模拟奖励导向的软最优策略，通过最小化模拟软最优策略与当前模型策略之间的 KL 散度来更新模型。我们的离策略建模结合 KL 散度最小化，相较于现有基于 RL 的方法在训练稳定性和样本效率方面具有优势。实验结果在蛋白质、小分子和调控 DNA 设计等多种任务中证明了该方法的有效性和优异的奖励优化能力。源代码已发布于 https://divelab.github.io/VIDD/。
引用
@article{arxiv.2507.00445,
title = {Iterative Distillation for Reward-Guided Fine-Tuning of Diffusion Models in Biomolecular Design},
author = {Xingyu Su and Xiner Li and Masatoshi Uehara and Sunwoo Kim and Yulai Zhao and Gabriele Scalia and Ehsan Hajiramezanali and Tommaso Biancalani and Degui Zhi and Shuiwang Ji},
journal= {arXiv preprint arXiv:2507.00445},
year = {2026}
}