DDPO — Training Diffusion Models with Reinforcement Learning

Diffusion models are a class of flexible generative models trained with an approximation to the log-likelihood objective. However, most use cases of diffusion models are not concerned with likelihoods, but instead with downstream objectives such as human-perceived image quality or drug effectiveness.

In this paper, we investigate reinforcement learning methods for directly optimizing diffusion models for such objectives. We describe how posing denoising as a multi-step decision-making problem enables a class of policy gradient algorithms, which we refer to as Denoising Diffusion Policy Optimization (DDPO), that are more effective than alternative reward-weighted likelihood approaches.

Empirically, DDPO can adapt text-to-image diffusion models to objectives that are difficult to express via prompting, such as image compressibility, and those derived from human feedback, such as aesthetic quality. Finally, we show that DDPO can improve prompt-image alignment using feedback from a vision-language model without the need for additional data collection or human annotation.

Diffusion probabilistic models have recently emerged as the de facto standard for generative modeling in continuous domains. Their flexibility in representing complex, high-dimensional distributions has led to the adoption of diffusion models in applications including image and video synthesis, drug and material design, and continuous control.

The key idea behind diffusion models is to iteratively transform a simple prior distribution into a target distribution by applying a sequential denoising process. This procedure is conventionally motivated as a maximum likelihood estimation problem, with the objective derived as a variational lower bound on the log-likelihood of the training data.

However, most use cases of diffusion models are not directly concerned with likelihoods, but instead with downstream objective such as human-perceived image quality or drug effectiveness. In this paper, we consider the problem of training diffusion models to satisfy such objectives directly, as opposed to matching a data distribution.

This problem is challenging because exact likelihood computation with diffusion models is intractable, making it difficult to apply many conventional reinforcement learning (RL) algorithms. We instead propose to frame denoising as a multi-step decision-making task, using the exact likelihoods at each denoising step in place of the approximate likelihoods induced by a full denoising process. We present a policy gradient algorithm, which we refer to as DDPO, that can optimize a diffusion model for downstream tasks using only a black-box reward function.

We apply our algorithm to the finetuning of large text-to-image diffusion models. Our initial evaluation focuses on tasks that are difficult to specify via prompting, such as image compressibility, and those derived from human feedback, such as aesthetic quality.

However, because many reward functions of interest are difficult to specify programmatically, finetuning procedures often rely on large-scale human labeling efforts to obtain a reward signal. In the case of text-to-image diffusion, we propose a method for replacing such labeling with feedback from a vision-language model (VLM). Similar to RLAIF finetuning for language models, the resulting procedure allows for diffusion models to be adapted to reward functions that would otherwise require additional human annotations. We use this procedure to improve prompt-image alignment for unusual subject-setting compositions.

Our contributions are as follows. We first present the derivation and conceptual motivation of DDPO. We then document the design of various reward functions for text-to-image generation, ranging from simple computations to workflows involving large VLMs, and demonstrate the effectiveness of DDPO compared to alternative reward-weighted likelihood methods in these settings. Finally, we demonstrate the generalization ability of our finetuning procedure to unseen prompts.

Diffusion probabilistic models. Denoising diffusion models have emerged as an effective class of generative models for modalities including images, videos, 3D shapes, and robotic trajectories. While the denoising objective is conventionally derived as an approximation to likelihood, the training of diffusion models typically departs from maximum likelihood in several ways. Modifying the objective to more strictly optimize likelihood often leads to worsened image quality, as likelihood is not a faithful proxy for visual quality. In this paper, we show how diffusion models can be optimized directly for downstream objectives.

Reinforcement learning from human feedback. A number of works have studied using human feedback to optimize models in settings such as simulated robotic control, game-playing, machine translation, citation retrieval, summarization, instruction-following, and alignment with specifications. Recently, Lee et al. studied the alignment of text-to-image diffusion models to human preferences using a method based on reward-weighted likelihood maximization. In our comparisons, their method corresponds to one iteration of the reward-weighted regression (RWR) method.

Diffusion models as sequential decision-making processes. Although predating diffusion models, Bachman et al. similarly posed data generation as a sequential decision-making problem and used the resulting framework to apply reinforcement learning methods to image generation. More recently, Fan et al. introduced a policy gradient method for training diffusion models. However, this paper aimed to improve data distribution matching rather than optimizing downstream objectives, and therefore the only reward function considered was a GAN-like discriminator.

In concurrent work to ours, DPOK built upon Fan et al. and Lee et al. to better align text-to-image diffusion models to human preferences using a policy gradient algorithm. Like Lee et al., DPOK only considers a single preference-based reward function; additionally, their work studies KL-regularization and primarily focuses on training a different diffusion model for each prompt. In contrast, we train on many prompts at once (up to 398) and demonstrate generalization to many more prompts outside the training set. Furthermore, we study how DDPO can be applied to multiple reward functions beyond those based on human feedback, including how rewards derived automatically from VLMs can improve prompt-image alignment.

In this work, we consider conditional diffusion probabilistic models, which represent a distribution $p(\mathbf{x}_0 \mid \mathbf{c})$ over a dataset of samples $\mathbf{x}_0$ and corresponding contexts $\mathbf{c}$. The distribution is modeled as the reverse of a Markovian forward process $q(\mathbf{x}_t \mid \mathbf{x}_{t-1})$, which iteratively adds noise to the data. Reversing the forward process can be accomplished by training a neural network $\boldsymbol{\mu}_\theta(\mathbf{x}_t, \mathbf{c}, t)$ with the following objective:

$$\mathcal{L}_{\text{diff}}(\theta) = \mathbb{E}_{(\mathbf{x}_0,\mathbf{c}),\,t\sim\mathcal{U}\{0,T\},\,\mathbf{x}_t\sim q}\left[\left\|\tilde{\boldsymbol{\mu}}(\mathbf{x}_0,t) - \boldsymbol{\mu}_\theta(\mathbf{x}_t,\mathbf{c},t)\right\|^2\right]$$

where $\tilde{\boldsymbol{\mu}}$ is the posterior mean of the forward process, a weighted average of $\mathbf{x}_0$ and $\mathbf{x}_t$. This objective is justified as maximizing a variational lower bound on the log-likelihood of the data.

Sampling begins with drawing a random $\mathbf{x}_T \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$ and following the reverse process $p_\theta(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{c})$. Most popular samplers use an isotropic Gaussian reverse process with a fixed timestep-dependent variance:

$$p_\theta(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{c}) = \mathcal{N}\!\left(\mathbf{x}_{t-1} \mid \boldsymbol{\mu}_\theta(\mathbf{x}_t,\mathbf{c},t),\ \sigma_t^2 \mathbf{I}\right)$$

A Markov decision process (MDP) is a formalization of sequential decision-making problems. An MDP is defined by a tuple $(\mathcal{S}, \mathcal{A}, \rho_0, P, R)$, in which $\mathcal{S}$ is the state space, $\mathcal{A}$ is the action space, $\rho_0$ is the distribution of initial states, $P$ is the transition kernel, and $R$ is the reward function.

At each timestep $t$, the agent observes a state $\mathbf{s}_t \in \mathcal{S}$, takes an action $\mathbf{a}_t \in \mathcal{A}$, receives a reward $R(\mathbf{s}_t, \mathbf{a}_t)$, and transitions to a new state $\mathbf{s}_{t+1} \sim P(\mathbf{s}_{t+1} \mid \mathbf{s}_t, \mathbf{a}_t)$. As the agent acts, it produces trajectories $\tau = (\mathbf{s}_0, \mathbf{a}_0, \dots, \mathbf{s}_T, \mathbf{a}_T)$. The RL objective is to maximize the expected cumulative reward:

$$\mathcal{J}_{\text{RL}}(\pi) = \mathbb{E}_{\tau \sim p(\tau \mid \pi)}\left[\sum_{t=0}^{T} R(\mathbf{s}_t, \mathbf{a}_t)\right]$$

We now describe how RL algorithms can be used to train diffusion models. We present two classes of methods and show that each corresponds to a different mapping of the denoising process to the MDP framework.

Problem statement. We assume a pre-existing diffusion model, which may be pretrained or randomly initialized. Assuming a fixed sampler, the diffusion model induces a sample distribution $p_\theta(\mathbf{x}_0 \mid \mathbf{c})$. The denoising diffusion RL objective is to maximize a reward signal $r$ defined on the samples and contexts:

$$\mathcal{J}_{\text{DDRL}}(\theta) = \mathbb{E}_{\mathbf{c}\sim p(\mathbf{c}),\ \mathbf{x}_0 \sim p_\theta(\mathbf{x}_0 \mid \mathbf{c})}\left[r(\mathbf{x}_0, \mathbf{c})\right]$$

for some context distribution $p(\mathbf{c})$ of our choosing.

Reward-weighted regression. To optimize $\mathcal{J}_{\text{DDRL}}$ with minimal changes to standard diffusion model training, we can use the denoising loss $\mathcal{L}_{\text{diff}}$, but with training data $\mathbf{x}_0 \sim p_\theta(\mathbf{x}_0 \mid \mathbf{c})$ and an added weighting that depends on the reward $r(\mathbf{x}_0, \mathbf{c})$. A standard weighting scheme uses exponentiated rewards to ensure nonnegativity,

$$w_{\text{RWR}}(\mathbf{x}_0, \mathbf{c}) = \frac{1}{Z}\exp\big(\beta r(\mathbf{x}_0, \mathbf{c})\big)$$

where $\beta$ is an inverse temperature and $Z$ a normalization constant. We also consider a simplified weighting scheme that uses binary weights, $w_{\text{sparse}}(\mathbf{x}_0,\mathbf{c}) = \mathbf{1}[r(\mathbf{x}_0,\mathbf{c}) \geq C]$, where $C$ is a reward threshold. In supervised learning terms, this is equivalent to repeated filtered finetuning on training data coming from the model.

Within the RL formalism, the RWR procedure corresponds to the following one-step MDP: $\mathbf{s} \triangleq \mathbf{c}$, $\mathbf{a} \triangleq \mathbf{x}_0$, $\pi(\mathbf{a}\mid\mathbf{s}) \triangleq p_\theta(\mathbf{x}_0\mid\mathbf{c})$, $\rho_0(\mathbf{s}) \triangleq p(\mathbf{c})$, $R(\mathbf{s},\mathbf{a}) \triangleq r(\mathbf{x}_0,\mathbf{c})$, with a transition kernel that immediately leads to an absorbing termination state.

From RL literature, weighting a log-likelihood objective by $w_{\text{RWR}}$ is known to approximately maximize $\mathcal{J}_{\text{RL}}(\pi)$ subject to a KL divergence constraint on $\pi$. However, $\mathcal{L}_{\text{diff}}$ does not involve an exact log-likelihood — it is instead derived as a variational bound on $\log p_\theta(\mathbf{x}_0 \mid \mathbf{c})$. Therefore, the RWR procedure applied to diffusion model training is not theoretically justified and only optimizes $\mathcal{J}_{\text{DDRL}}$ very approximately.

Denoising diffusion policy optimization. RWR relies on an approximate log-likelihood because it ignores the sequential nature of the denoising process, only using the final samples $\mathbf{x}_0$. In this section, we show how the denoising process can be reframed as a multi-step MDP, allowing us to directly optimize $\mathcal{J}_{\text{DDRL}}$ using policy gradient estimators.

Denoising as a multi-step MDP. We map the iterative denoising procedure to the following MDP:

$$\mathbf{s}_t \triangleq (\mathbf{c}, t, \mathbf{x}_t) \qquad \mathbf{a}_t \triangleq \mathbf{x}_{t-1} \qquad \pi(\mathbf{a}_t \mid \mathbf{s}_t) \triangleq p_\theta(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{c})$$

$$P(\mathbf{s}_{t+1} \mid \mathbf{s}_t, \mathbf{a}_t) \triangleq \big(\delta_{\mathbf{c}}, \delta_{t-1}, \delta_{\mathbf{x}_{t-1}}\big) \qquad \rho_0(\mathbf{s}_0) \triangleq \big(p(\mathbf{c}), \delta_T, \mathcal{N}(\mathbf{0}, \mathbf{I})\big)$$

$$R(\mathbf{s}_t, \mathbf{a}_t) \triangleq \begin{cases} r(\mathbf{x}_0, \mathbf{c}) & \text{if } t = 0 \\ 0 & \text{otherwise} \end{cases}$$

in which $\delta_y$ is the Dirac delta distribution with nonzero density only at $y$. Trajectories consist of $T$ timesteps, after which $P$ leads to a termination state. The cumulative reward of each trajectory is equal to $r(\mathbf{x}_0, \mathbf{c})$, so maximizing $\mathcal{J}_{\text{DDRL}}(\theta)$ is equivalent to maximizing $\mathcal{J}_{\text{RL}}(\pi)$ in this MDP.

The benefit of this formulation is that if we use a standard sampler with $p_\theta(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{c})$ parameterized as in Equation 2, the policy $\pi$ becomes an isotropic Gaussian as opposed to the arbitrarily complicated distribution $p_\theta(\mathbf{x}_0 \mid \mathbf{c})$ as it is in the RWR formulation. This simplification allows for the evaluation of exact log-likelihoods and their gradients with respect to the diffusion model parameters.

Policy gradient estimation. With access to likelihoods and likelihood gradients, we can make direct Monte Carlo estimates of $\nabla_\theta \mathcal{J}_{\text{DDRL}}$. Like RWR, DDPO alternates collecting denoising trajectories $\{\mathbf{x}_T, \mathbf{x}_{T-1}, \dots, \mathbf{x}_0\}$ via sampling and updating parameters via gradient descent.

The first variant of DDPO, which we call DDPO$_{\text{SF}}$, uses the score function policy gradient estimator, also known as the likelihood ratio method or REINFORCE:

$$\nabla_\theta \mathcal{J}_{\text{DDRL}} = \mathbb{E}\left[\sum_{t=0}^{T} \nabla_\theta \log p_\theta(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{c}) \; r(\mathbf{x}_0, \mathbf{c})\right]$$

where the expectation is taken over denoising trajectories generated by the current parameters $\theta$. However, this estimator only allows for one step of optimization per round of data collection, as the gradient must be computed using data generated by the current parameters. To perform multiple steps of optimization, we may use an importance sampling estimator:

$$\nabla_\theta \mathcal{J}_{\text{DDRL}} = \mathbb{E}\left[\sum_{t=0}^{T} \frac{p_\theta(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{c})}{p_{\theta_{\text{old}}}(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{c})} \nabla_\theta \log p_\theta(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{c}) \; r(\mathbf{x}_0, \mathbf{c})\right]$$

which we refer to as DDPO$_{\text{IS}}$. This estimator becomes inaccurate if $p_\theta$ deviates too far from $p_{\theta_{\text{old}}}$, which can be addressed using trust regions to constrain the size of the update. In practice, we implement the trust region via clipping, as in proximal policy optimization (PPO).

In this work, we evaluate our methods on text-to-image diffusion. Text-to-image diffusion serves as a valuable test environment for reinforcement learning due to the availability of large pretrained models and the versatility of using diverse and visually interesting reward functions. We study a spectrum of reward functions of varying complexity, ranging from those that are straightforward to specify and evaluate to those that capture the depth of real-world downstream tasks.

Compressibility and incompressibility. The capabilities of text-to-image diffusion models are limited by the co-occurrences of text and images in their training distribution. For instance, images are rarely captioned with their file size, making it impossible to specify a desired file size via prompting. This limitation makes reward functions based on file size a convenient case study: they are simple to compute, but not controllable through the conventional methods of likelihood maximization and prompt engineering.

We fix the resolution of diffusion model samples at 512×512, such that the file size is determined solely by the compressibility of the image. We define two tasks based on file size: compressibility, in which the file size of the image after JPEG compression is minimized, and incompressibility, in which the same measure is maximized.

Aesthetic quality. To capture a reward function that would be useful to a human user, we define a task based on perceived aesthetic quality. We use the LAION aesthetics predictor, which is trained on 176,000 human image ratings. The predictor is implemented as a linear model on top of CLIP embeddings. Annotations range between 1 and 10, with the highest-rated images mostly containing artwork.

Since the aesthetic quality predictor is trained on human judgments, this task constitutes reinforcement learning from human feedback.

Automated prompt alignment with vision-language models. A very general-purpose reward function for training a text-to-image model is prompt-image alignment. However, specifying a reward that captures generic prompt alignment is difficult, conventionally requiring large-scale human labeling efforts. We propose using an existing VLM to replace additional human annotation. This design is inspired by recent work on RLAIF, in which language models are improved using feedback from themselves.

We use LLaVA, a state-of-the-art VLM, to describe an image. The finetuning reward is the BERTScore recall metric, a measure of semantic similarity, using the prompt as the reference sentence and the VLM description as the candidate sentence. Samples that more faithfully include all of the details of the prompt receive higher rewards, to the extent that those visual details are legible to the VLM.

In Figure 4, we show one simple question: “what is happening in this image?”. While this captures the general task of prompt-image alignment, in principle any question could be used to specify complex or hard-to-define reward functions for a particular use case. One could even employ a language model to automatically generate candidate questions and evaluate responses based on the prompt. This framework provides a flexible interface where the complexity of the reward function is only limited by the capabilities of the vision and language models involved.

The purpose of our experiments is to evaluate the effectiveness of RL algorithms for finetuning diffusion models to align with a variety of user-specified objectives. After examining the viability of the general approach, we focus on the following questions:

(1) How do variants of DDPO compare to RWR and to each other? (2) Can VLMs allow for optimizing rewards that are difficult to specify manually? (3) Do the effects of RL finetuning generalize to prompts not seen during finetuning?

Algorithm comparisons. We begin by evaluating all methods on the compressibility, incompressibility, and aesthetic quality tasks, as these tasks isolate the effectiveness of the RL approach from considerations relating to the VLM reward function. We use Stable Diffusion v1.4 as the base model for all experiments. Compressibility and incompressibility prompts are sampled uniformly from all 398 animals in the ImageNet-1000 categories. Aesthetic quality prompts are sampled uniformly from a smaller set of 45 common animals.

As shown qualitatively in Figure 2, DDPO is able to effectively adapt a pretrained model with only the specification of a reward function and without any further data curation. The strategies found to optimize each reward are nontrivial; for example, to maximize LAION-predicted aesthetic quality, DDPO transforms a model that produces naturalistic images into one that produces artistic drawings. To maximize compressibility, DDPO removes backgrounds and applies smoothing to what remains. To maximize incompressibility, DDPO finds artifacts that are difficult for the JPEG compression algorithm to encode, such as high-frequency noise and sharp edges.

We provide a quantitative comparison of all methods in Figure 3. We plot the attained reward as a function of the number of queries to the reward function, as reward evaluation becomes the limiting factor in many practical applications. DDPO shows a clear advantage over RWR on all tasks, demonstrating that formulating the denoising process as a multi-step MDP and estimating the policy gradient directly is more effective than optimizing a reward-weighted variational bound on log-likelihood. Within the DDPO class, the importance sampling estimator slightly outperforms the score function estimator, likely due to the increased number of optimization steps. Within the RWR class, the performance of weighting schemes is comparable, making the sparse weighting scheme preferable on these tasks due to its simplicity and reduced resource requirements.

Automated prompt alignment. We next evaluate the ability of VLMs, in conjunction with DDPO, to automatically improve the image-prompt alignment of the pretrained model without additional human labels. We focus on DDPO$_{\text{IS}}$ for this experiment, as we found it to be the most effective algorithm. The prompts for this task all have the form “a(n) [animal] [activity]”, where the animal comes from the same list of 45 common animals and the activity is chosen from a list of 3 activities: “riding a bike”, “playing chess”, and “washing dishes”.

Qualitatively, the samples come to depict the prompts much more faithfully throughout the course of training. This trend is also reflected quantitatively, though is less salient as small changes in BERTScore can correspond to large differences in relevance. It is important to note that some of the prompts in the finetuning set, such as “a dolphin riding a bike”, had zero success rate from the pretrained model; if trained in isolation, this prompt would be unlikely to ever improve because there would be no reward signal. It was only via transferrable learning across prompts that these difficult prompts could improve.

Nearly all of the samples become more cartoon-like or artistic during finetuning. This was not optimized for directly. We hypothesize that this may be a function of the pretraining distribution or of the reward function (perhaps LLaVA has an easier time recognizing the content of simple cartoon-like images).

Generalization. RL finetuning on large language models has been shown to produce interesting generalization properties; for example, instruction finetuning almost entirely in English has been shown to improve capabilities in other languages. It is difficult to reconcile this phenomenon with our current understanding of generalization; it would a priori seem more likely for finetuning to have an effect only on the finetuning prompt set or distribution.

In order to investigate the same phenomenon with diffusion models, Figure 6 shows a set of DDPO-finetuned model samples corresponding to prompts that were not seen during finetuning. In concordance with instruction-following transfer in language modeling, we find that the effects of finetuning do generalize, even with prompt distributions as narrow as 45 animals and 3 activities. We find evidence of generalization to animals outside of the training distribution, to non-animal everyday objects, and in the case of prompt-image alignment, even to novel activities such as “taking an exam”.

We presented an RL-based framework for training denoising diffusion models to directly optimize a variety of reward functions. By posing the iterative denoising procedure as a multi-step decision-making problem, we were able to design a class of policy gradient algorithms that are highly effective at training diffusion models. We found that DDPO was an effective optimizer for tasks that are difficult to specify via prompts, such as image compressibility, and difficult to evaluate programmatically, such as semantic alignment with prompts.

While our evaluation considers a variety of prompts, the full range of images in our experiments was constrained (e.g., animals performing activities). Future iterations could expand both the questions posed to the VLM, possibly using language models to propose relevant questions based on the prompt, as well as the diversity of the prompt distribution. We also chose not to study the problem of overoptimization, a common issue with RL finetuning in which the model diverges too far from the original distribution to be useful; we highlight this as an important area for future work.

We hope that this work will provide a step toward more targeted training of large generative models, where optimization via RL can produce models that are effective at achieving user-specified goals rather than simply matching an entire data distribution.

Broader Impacts. Generative models can be valuable productivity aids, but may also pose harm when used for disinformation, impersonation, or phishing. Our work aims to make diffusion models more useful by enabling them to optimize user-specified objectives. This adaptation has beneficial applications, such as the generation of more understandable educational material, but may also be used maliciously. Work on the reliable detection of synthetic content remains important to mitigate such harms.