[论文解读] Understanding Reinforcement Learning-Based Fine-Tuning of Diffusion Models: A Tutorial and Review

Masatoshi Uehara, Yulai Zhao+2 more|arXiv (Cornell University)|Jul 18, 2024

Traffic control and management被引用 4

一句话总结

本教程提出了一种基于强化学习（RL）的统一框架，用于微调预训练扩散模型，以优化下游奖励函数，例如蛋白质稳定性或分子活性。通过将去噪过程形式化为马尔可夫决策过程（MDP），该框架使PPO、奖励加权MLE和路径一致性学习等RL算法能够直接最大化期望指标，同时保持样本质量和泛化能力。

ABSTRACT

This tutorial provides a comprehensive survey of methods for fine-tuning diffusion models to optimize downstream reward functions. While diffusion models are widely known to provide excellent generative modeling capability, practical applications in domains such as biology require generating samples that maximize some desired metric (e.g., translation efficiency in RNA, docking score in molecules, stability in protein). In these cases, the diffusion model can be optimized not only to generate realistic samples but also to explicitly maximize the measure of interest. Such methods are based on concepts from reinforcement learning (RL). We explain the application of various RL algorithms, including PPO, differentiable optimization, reward-weighted MLE, value-weighted sampling, and path consistency learning, tailored specifically for fine-tuning diffusion models. We aim to explore fundamental aspects such as the strengths and limitations of different RL-based fine-tuning algorithms across various scenarios, the benefits of RL-based fine-tuning compared to non-RL-based approaches, and the formal objectives of RL-based fine-tuning (target distributions). Additionally, we aim to examine their connections with related topics such as classifier guidance, Gflownets, flow-based diffusion models, path integral control theory, and sampling from unnormalized distributions such as MCMC. The code of this tutorial is available at https://github.com/masa-ue/RLfinetuning_Diffusion_Bioseq

研究动机与目标

• 将基于强化学习的扩散模型微调形式化为熵正则化的MDP中的序列决策问题。
• 对比和分析多种RL算法（如PPO、奖励加权MLE、价值加权采样和路径一致性学习）在微调扩散模型方面的性能。
• 根据奖励反馈的可用性（在线与离线、已知与未知奖励）对微调场景进行分类，并分析其对算法选择的影响。
• 明确基于强化学习的微调形式目标，包括离线学习中目标分布与分布偏移风险的关系。
• 建立基于强化学习的微调与相关概念（如分类器引导、GflowNets、基于流的模型以及路径积分控制理论）之间的联系。

提出的方法

• 将扩散去噪过程形式化为马尔可夫决策过程（MDP），其中每个去噪步骤均作为为最大化累积奖励而做出的决策。
• 应用现成的RL算法，如近端策略优化（PPO）、可微分奖励反向传播和奖励加权最大似然估计（MLE），以优化策略。
• 采用价值加权采样和路径一致性学习（PCL）以提升样本质量与训练稳定性，尤其适用于高维和多模态分布。
• 将预训练的扩散模型用作参考SDE（随机微分方程），避免从零开始学习先验分布，从而实现更快的推理速度。
• 通过分析离线RL中分布偏移的风险，包括对分布外样本的过度优化问题，并提出缓解策略。
• 整合与相关框架的联系：分类器引导（作为价值加权采样的特例）、GflowNets（用于基于流的建模），以及MCMC在未归一化分布采样中的应用。

Figure 1 : Illustrative examples of RL-based fine-tuning, aimed at optimizing pre-trained diffusion models to maximize downstream reward functions.

实验结果

研究问题

• RQ1 扩散模型中的序列去噪过程如何自然地被建模为强化学习问题？
• RQ2 在微调扩散模型以实现下游奖励最大化方面，不同RL算法（如PPO、奖励加权MLE、PCL）的相对优势与局限性是什么？
• RQ3 反馈获取场景（在线、离线、已知与未知奖励）如何影响基于强化学习的微调中的算法设计与性能表现？
• RQ4 基于强化学习的微调的正式目标是什么？它与目标分布及分布偏移之间有何关系？
• RQ5 基于强化学习的微调方法与分类器引导、GflowNets以及路径积分控制理论等既有概念之间存在何种联系？

主要发现

• 基于强化学习的微调可直接优化扩散模型以实现下游奖励函数（如蛋白质稳定性、RNA翻译效率和分子QED评分）的提升，在目标生成方面优于非RL基线方法。
• 奖励加权MLE和PPO作为有效基线，其中PPO在高维、多模态设置中表现优异，得益于其稳定性和策略优化框架。
• 路径一致性学习（PCL）与价值加权采样可提供稳定的训练动态并减少模式崩溃，尤其在复杂奖励景观中表现突出。
• 将预训练扩散模型用作参考SDE可显著缩短训练时间，并避免估计先验分布的需要，相较基于MCMC的方法更具优势。
• 在离线RL设置中，微调模型容易因反馈覆盖有限而对分布外样本产生过度优化，凸显了对分布鲁棒性需求的重要性。
• 该框架建立了基于强化学习的微调与分类器引导之间的正式联系，表明分类器引导是特定奖励设计下价值加权采样的特例。

Figure 4 : Offline scenario