[Paper Review] 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 management4 citations

TL;DR

This tutorial presents a unified reinforcement learning (RL)-based framework for fine-tuning pre-trained diffusion models to optimize downstream reward functions, such as protein stability or molecular activity. By formulating the denoising process as a Markov decision process, it enables the use of RL algorithms like PPO, reward-weighted MLE, and path consistency learning to directly maximize desired metrics while maintaining sample quality and generalization.

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

Motivation & Objective

• To formalize RL-based fine-tuning of diffusion models as a sequential decision-making problem in entropy-regularized MDPs.
• To compare and contrast various RL algorithms—PPO, reward-weighted MLE, value-weighted sampling, and path consistency learning—for fine-tuning diffusion models.
• To categorize fine-tuning scenarios based on reward feedback availability (online vs. offline, known vs. unknown rewards) and analyze implications for algorithm selection.
• To clarify the formal objectives of RL-based fine-tuning, including the target distribution and distributional shift risks in offline learning.
• To establish connections between RL-based fine-tuning and related concepts such as classifier guidance, GflowNets, flow-based models, and path integral control theory.

Proposed method

• Formalize the diffusion denoising process as a Markov decision process (MDP), where each denoising step is a decision to maximize cumulative reward.
• Apply off-the-shelf RL algorithms such as Proximal Policy Optimization (PPO), differentiable reward backpropagation, and reward-weighted maximum likelihood estimation (MLE) to optimize the policy.
• Use value-weighted sampling and path consistency learning (PCL) to improve sample quality and training stability, especially in high-dimensional and multimodal distributions.
• Leverage pre-trained diffusion models as the reference SDE (stochastic differential equation), avoiding the need to learn a prior from scratch and enabling faster inference.
• Address distributional shift in offline RL by analyzing risks of overoptimization on out-of-distribution samples and proposing mitigation strategies.
• Integrate connections with related frameworks: classifier guidance (as a special case of value-weighted sampling), GflowNets (for flow-based modeling), and MCMC for unnormalized distribution sampling.

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

Experimental results

Research questions

• RQ1 How can the sequential denoising process in diffusion models be naturally framed as a reinforcement learning problem?
• RQ2 What are the relative strengths and limitations of different RL algorithms (e.g., PPO, reward-weighted MLE, PCL) in fine-tuning diffusion models for downstream reward maximization?
• RQ3 How do feedback acquisition scenarios (online, offline, known vs. unknown rewards) influence algorithm design and performance in RL-based fine-tuning?
• RQ4 What is the formal objective of RL-based fine-tuning, and how does it relate to the target distribution and distributional shift?
• RQ5 How are RL-based fine-tuning methods connected to established concepts like classifier guidance, GflowNets, and path integral control theory?

Key findings

• RL-based fine-tuning enables direct optimization of diffusion models for downstream reward functions such as protein stability, RNA translation efficiency, and molecular QED scores, outperforming non-RL baselines in targeted generation.
• Reward-weighted MLE and PPO are effective baselines, with PPO showing strong performance in high-dimensional, multimodal settings due to its stability and policy optimization framework.
• Path consistency learning (PCL) and value-weighted sampling provide stable training dynamics and reduce mode collapse, particularly in complex reward landscapes.
• Using pre-trained diffusion models as the reference SDE significantly reduces training time and avoids the need to estimate the prior distribution, unlike MCMC-based methods.
• In offline RL settings, fine-tuned models are prone to overoptimization on out-of-distribution samples due to limited feedback coverage, highlighting the need for distributional robustness.
• The framework establishes formal connections between RL-based fine-tuning and classifier guidance, showing that classifier guidance is a special case of value-weighted sampling with specific reward shaping.

Figure 4 : Offline scenario