Adding Conditional Control to Diffusion Models with Reinforcement Learning

## Abstract

Abstract Diffusion models are powerful generative models that allow for precise control over the characteristics of the generated samples. While these diffusion models trained on large datasets have achieved success, there is often a need to introduce additional controls in downstream fine-tuning processes, treating these powerful models as pre-trained diffusion models. This work presents a novel method based on reinforcement learning (RL) to add additional controls, leveraging an offline dataset comprising inputs and corresponding labels. We formulate this task as an RL problem, with the classifier learned from the offline dataset and the KL divergence against pre-trained models serving as the reward functions. We introduce our method, CTRL ( C onditioning pre- T rained diffusion models with R einforcement L earning), which produces soft-optimal policies that maximize the abovementioned reward functions. We formally demonstrate that our method enables sampling from the conditional distribution conditioned on additional controls during inference.
Our RL-based approach offers several advantages over existing methods. Compared to commonly used classifier-free guidance,
our approach improves sample efficiency, and can greatly simplify offline dataset construction by exploiting conditional independence between the inputs and additional controls. Furthermore, unlike classifier guidance, we avoid the need to train classifiers from intermediate states to additional controls.

## 1 Introduction

Diffusion models have emerged as effective generative models for capturing intricate distributions (Sohl-Dickstein et al., ; Ho et al., ). Their capabilities are further enhanced by building conditional diffusion models $p(x|c)$. For instance, in text-to-image generative models like DALL-E (Ramesh et al., ) and Stable Diffusion (Rombach et al., ), $c\in\mathcal{C}$ is a prompt, and $x\in\mathcal{X}$ is the image generated according to this prompt.
While diffusion models trained on extensive datasets have shown remarkable success, additional controls often need to be incorporated during the downstream fine-tuning process when treating these powerful models as pre-trained diffusion models.

In this work, our goal is to incorporate new conditional controls into pre-trained diffusion models. Specifically, given access to a large pre-trained model capable of modeling $p(x|c)$ trained on extensive datasets, we aim to condition it on an additional random variable $y\in\mathcal{Y}$, thereby creating a generative model $p(x|c,y)$. To accomplish this, we utilize the pre-trained model and an offline dataset consisting of triplets of $\{c,x,y\}$. This scenario is important, as highlighted in the existing literature on computer vision (e.g., Zhang et al. ()), because it enables the extension of generative capabilities with new conditional variables without requiring retraining from scratch.
Currently, classifier-free guidance (Ho et al., ) is a prevailing approach for incorporating conditional controls into diffusion models, and it has proven successful in computer vision  (Zhang et al., ; Zhao et al., ). However, its effectiveness may not be generalized well to specialized domains (e.g., scientific domains) where collecting large offline datasets is prohibitively expensive. Indeed, the success of training conditional diffusion models via classifier-free guidance heavily depends on the availability of sufficiently large offline datasets (Brooks et al., ), which is often not feasible in many scenarios, e.g. drug discovery (Huang et al., ).

In our work, we present a new approach for adding new conditional controls
via reinforcement learning (RL) to further improve sample efficiency. Inspired by recent progress in RL-based fine-tuning (Black et al., ; Fan et al., ), we frame the conditional generation as an RL problem within a Markov Decision Process (MDP). In this formulation, the reward, which we want to maximize, is defined as the (conditional) log-likelihood function $\log p(y|x,c)$, and the policy, which is dependent on $(c,y)$, corresponds to the denoising process at each time step in a diffusion model. We formally demonstrate that, by executing the soft-optimal policy, which maximizes the reward $\log p(y|x,c)$ with KL penalty against the pre-trained model, we can sample from the target conditional distribution $p(x|c,y)$ during inference. Hence, our proposed algorithm, CTRL (Conditioning pre-Trained diffusion models
with Reinforcement Learning) consists of three main steps: (1) learning a classifier $\log p(y|x,c)$ (which will serve as our reward function in the MDP) from the offline dataset, (2) constructing an augmented diffusion model by adding (trainable) parameters to the pre-trained model in order to accommodate an additional label $y$, and (3) learning soft-optimal policy within the aforementioned MDP during fine-tuning. Our approach is novel as it significantly diverges from classifier-free guidance and distinguishes itself from existing RL-based fine-tuning methods by integrating an augmented model in the fine-tuning process to support additional controls.

**Table 1: Comparison between our proposal and existing approaches. In contrast to classifier guidance or its variations, our method entails re-training the models directly on top of pre-trained models (i.e., fine-tuning). Additionally, we circumvent the necessity of learning a mapping $x_{t}\to y$ or employing heuristic approximation techniques to address this issue. Compared to classifier-free guidance which always demands triplets $\{c,x,y\}$, our method can leverage conditional independence and only necessitate pairs $\{x,y\}$ by leveraging if $y\perp c|x$ holds. This simplifies the construction of the offline dataset.**
| Methods | Fine-tuning | Need to learn $x_{t}\to y$ | Leveraging conditional independence |
| --- | --- | --- | --- |
| Classifier guidance (Dhariwal and Nichol, ) | No | Yes | Yes |
| Reconstruction guidance (e.g.  (Ho et al., ),  (Chung et al., ),<br>(Han et al., )) | No | No | Yes |
| Classifier-free guidance (Ho and Salimans, ) | Yes | No | No |
| CTRL (Ours) | Yes | No | Yes |

Our novel RL-based approach offers several advantages over existing methods for adding additional controls. Firstly, in contrast to classifier-free guidance, which uses offline data to directly model $p(x|y,c)$, our method effectively leverages offline data by modeling a simpler distribution $p(y|x,c)$ (in typical scenarios where $y$ is lower dimensional than $x$).
This enhances the sample efficiency of our approach. Secondly, in typical scenarios where an additional label $y$ depends only on $x$ (for example, the compressibility of images only concerns images, not the prompts), our fine-tuning method only requires $\{x,y\}$ pairs, whereas classifier-free guidance still necessitates triplets $(c,x,y)$ from the offline dataset. This is due to the fact that the reward function is simplified to $\log p(y|x)$ because the conditional independence $y\perp c|x$ leads to $\log p(y|x,c)=\log p(y|x)$. Moreover, when the goal is to simultaneously add conditioning controls on two labels, $y_{1}$ and $y_{2}$, and both labels only depend on $x$, our method solely requires $\{x,y_{1}\}$ pairs and $\{x,y_{2}\}$ pairs, while classifier-free guidance requires quadruples $\{c,x,y_{1},y_{2}\}$. Therefore, in this manner, CTRL can also exploit the *compositional nature* of the mapping between inputs and additional labels.

Our contributions can be summarized as follows. We propose an RL-based fine-tuning approach for conditioning pre-trained diffusion models on additional labels. In comparison to classifier-free guidance, the proposed method uses the offline dataset in a sample-efficient manner and enables leveraging the conditional independence assumption, which significantly simplifies the construction of the offline dataset. Also, we establish a close connection with classifier guidance (Dhariwal and Nichol, ; Song et al., ) by proving that it is actually another way to obtain the abovementioned soft-optimal policies (in ideal cases where there are no statistical/model-misspecification errors in algorithms). Despite this similarity, our algorithm addresses common challenges associated with classifier guidance, such as the need to learn classifiers at multiple noise scales, or certain fundamental approximations to overcome these challenges (Chung et al., ; Song et al., ). Table 1 summarizes the main features of the proposed algorithm compared to existing methods.

## 2 Related Works

##### Classfier guidance.

Dhariwal and Nichol (); Song et al. ()
introduced classifier guidance, a method that entails training a classifier and incorporating its gradients to guide inference (while freezing pre-trained models). However, a notable drawback of this technique lies in the classifier’s accuracy in predicting $y$ from intermediate $x_{t}$, resulting in cumulative errors during the diffusion process. To address this issue, several studies propose methods to circumvent it by employing certain approximations, using a mapping from intermediate states $x_{t}$ to the original input space $x_{0}$, and solely learning a classifier from $x_{0}$ to $y$ (Ho et al., ; Han et al., ; Chung et al., ; Finzi et al., ; Bansal et al., ). In contrast to these works, our approach focuses on fine-tuning the diffusion model itself rather than relying on an inference-time technique. While the strict comparison between model fine-tuning and inference-time techniques is not feasible, we theoretically elucidate the distinctions and connections of our approach with classifier guidance in Section 5.1.

##### Classfier-free guidance.

Classifier-free guidance (Ho and Salimans, ) is a method that directly conditions the generative process on both data and context, bypassing the need for explicit classifiers. This methodology has been widely and effectively applied, for example, in text-to-image models (Nichol et al., ; Saharia et al., ; Rombach et al., ). While the original research does not explore classifier-free guidance within the scope of fine-tuning pre-trained diffusion models, several subsequent studies address fine-tuning scenarios Zhang et al. (); Xie et al. (). As elucidated in Section 5.2, compared to classifier-free guidance, our approach can improve sample efficiency and leverage conditional independence to facilitate the construction of the offline dataset used for fine-tuning.

##### Fine-tuning via RL.

Several previous studies have addressed the fine-tuning of diffusion models by optimizing relevant reward functions. Methodologically, these approaches encompass supervised learning (Lee et al., ; Wu et al., ), reinforcement learning (Black et al., ; Fan et al., ; Uehara et al., ), and control-based techniques (Clark et al., ; Xu et al., ; Prabhudesai et al., ; Uehara et al., ). While our proposal draws inspiration from these works, our objective for fine-tuning is to tackle a distinct goal: incorporating *additional* controls. To achieve this, unlike previous approaches, we employ policies with augmented parameters, rather than merely fine-tuning pre-trained models without adding any new parameters.

## 3 Preliminaries

In this section, we introduce the problem setting, review the existing methods addressing this problem, and discuss their disadvantages.

### 3.1 Goal: Conditioning with Additional Labels Using Offline Data

We first define our main setting and main objective.
Throughout this paper, we use $\mathcal{Y}$ and $\mathcal{C}$ to represent condition spaces and $\mathcal{X}$ to denote the (Euclidean) sample space. Given the pre-trained model, which enables us to sample from $p^{\mathrm{pre}}(x|c):\mathcal{C}\to\Delta(\mathcal{X})$, our goal is to add new conditional controls $y\in\mathcal{Y}$ such that we can sample from $p(x|c,y)$.

##### Pre-trained model and offline dataset.

A (continuous-time) pre-trained conditional diffusion model is characterized by the following SDE(^1^11To simplify the notation, as in the case of bridge matching, we present a case when an initial distribution is Dirac delta distribution. In the case of stochastic distribution, the extension of our proposal is still straightforward, as done in (Uehara et al., ).):

$$ $\displaystyle dx_{t}=f^{\mathrm{pre}}(t,c,x_{t};\theta^{\mathrm{pre}})dt+ \sigma(t)dw_{t},\quad x_{0}=x_{\mathrm{ini}},$ (1) $$

where $f^{\mathrm{pre}}:[0,T]\times\mathcal{C}\times\mathcal{X}\to\mathbb{R}^{d}$ is a model with parameter $\theta$.

In training diffusion models, the parameter $\theta^{\mathrm{pre}}$ is derived by optimizing a specific loss function on large datasets(^2^22For notational simplicity, throughout this work, we would often drop $\theta^{\mathrm{pre}}$.). We refer interested readers to  Appendix A for more details on how to construct these loss functions. Using the pre-trained model and following the above SDE (1) from 0 0 to $T$, we can sample from $p^{\mathrm{pre}}(\cdot|c)$ for any condition $c\in\mathcal{C}$.

In addition to the pre-trained model, to add additional control, as in many recent works (Dhariwal and Nichol, ; Bansal et al., ; Epstein et al., ), we assume that we have access to offline data: $\mathcal{D}=\{c^{(i)},x^{(i)},y^{(i)}\}_{i=1}^{n}\in\mathcal{C}\times\mathcal{
X}\times\mathcal{Y}$. We denote the conditional distribution of $y$ given $x$ and $c$ by $p^{\diamond}(y|x,c)$.

##### Target distribution.

Using the pre-trained model and the offline dataset, our goal is to obtain a diffusion model such that we can sample from a distribution over $\mathcal{C}\times\mathcal{Y}\to\Delta(\mathcal{X})$ as below:

$$ $\displaystyle p_{\gamma}(\cdot|c,y):=\frac{\{p^{\diamond}(y|\cdot,c)\}^{\gamma }p^{\mathrm{pre}}(\cdot|c)}{\int\{p^{\diamond}(y|x,c)\}^{\gamma}p^{\mathrm{pre }}(x|c)\mu(dx)},$ (2) $$

where the parameter $\gamma\in\mathbb{R}^{+}$ represents the strength of the additional guidance and $\mu$ is the Lebsgue measure.

Such target distribution is extensively explored in the literature on classifier guidance and classifier-free guidance (Dhariwal and Nichol, ; Ho and Salimans, ; Nichol et al., ; Saharia et al., ; Rombach et al., ). Specifically, when $\gamma=1$, this distribution corresponds to the standard conditional distribution $p(x|c,y)$, which is a fundamental objective of many conditional generative models (Dhariwal and Nichol, ; Ho and Salimans, ). Moreover, for a general $\gamma$, $p_{\gamma}$ can be formulated via the following optimization problem:

$$ $\displaystyle p_{\gamma}(\cdot|c,y)=\mathop{\mathrm{argmin}}_{q:\mathcal{C} \times\mathcal{Y}\to\Delta(\mathcal{X})}\mathbb{E}_{x\sim q(\cdot|c,y)}[- \gamma\log p^{\diamond}(y|x,c)]+\mathrm{KL}(q(\cdot|c,y)\|p^{\mathrm{pre}}( \cdot|c)).$ $$

This relation is clear by noticing that the objective function is equal to $\mathrm{KL}(q(\cdot|c,y)\|p_{\gamma}(\cdot|c,y))$ up to a constant.

##### Goal.

As discussed, the primary goal of this research is to train a generative model capable of simulating $p_{\gamma}(\cdot|c,y)$. To achieve this, we introduce the following SDE:

$$ $\displaystyle dx_{t}=g(t,c,y,x_{t})\,dt+\sigma(t)\,dw_{t},\quad x_{0}=x_{\text {ini}},$ (3) $$

where $g:[0,T]\times\mathcal{C}\times\mathcal{Y}\times\mathcal{X}\to\mathbb{R}^{d}$ is an augmented model to add additional controls into pre-trained models. The primary challenge involves leveraging both offline data and pre-trained model weights to train the term $g$, ensuring that the marginal distribution of $x_{T}$ induced by the SDE (3) accurately approximates $p_{\gamma}$.

##### Notation.

Let the space of trajectories $x_{0:T}$ be $\mathcal{K}$. Conditional on $c$ and $y$, we denote the measure induced by the SDE (3) over $\mathcal{K}$ by $\mathbb{P}^{g}(\cdot|c,y)$. Similarly, we use $\mathbb{P}^{g}_{t}(\cdot|c,y)$ and $p^{g}_{t}(\cdot|c,y)$ to represent the marginal distribution of $x_{t}$ and density $d\mathbb{P}^{g}_{t}(\tau|c,y)/d\mu$.

###### Remark 1 (Extension to the non-Euclidean space) .

To streamline the notation, we focus on scenarios where $\mathcal{X}$ is Euclidian.
However, when dealing with discrete spaces, we can still extend our discussion by examining the discretized version from the beginning (Uehara et al., , Theorem 1).

###### Remark 1 (Extension to the non-Euclidean space) .

### 3.2 Existing Methods

Here we elucidate the existing methods for conditional generation. Two approaches are available to achieve the aforementioned goal by harnessing the pre-trained model.

#### 3.2.1 Classfier-Free Guidance

While the original work (Ho and Salimans, ) did not incorporate a pre-trained model, subsequent studies have explored the extension of classifier-free guidance to pre-trained models (e.g., InstructPix2Pix (Brooks et al., ),
ControlNet (Zhang et al., ), DiffFit (Xie et al., )). These methods introduce an augmented model $g$ as described in (3), where the weights are initialized from the pre-trained model. Then, they fine-tune by minimizing the classifier-free guidance loss on the offline dataset. Although these approaches have demonstrated success in computer vision, their effectiveness may not translate to other fields, such as scientific domains, where the size of the offline datasets for new conditions is often limited (Huang et al., ; Yellapragada et al., ; Giannone et al., ).

#### 3.2.2 Classfier Guidance

Classifier guidance (Dhariwal and Nichol, ; Song et al., ) is based on the following celebrated result.

###### Lemma 1 (Doob’s h-transforms (Rogers and Williams, 2000 ) ) .

For any $c\in\mathcal{C}$ and $y\in\mathcal{Y}$, by evolving according to the following SDE from 0 0 to $T$:

$$ $\displaystyle dx_{t}=\{f^{\mathrm{pre}}(t,c,x_{t})+\sigma^{2}(t)\underbrace{ \nabla_{x_{t}}\log\mathbb{E}_{x_{t:T}\sim\mathbb{P}^{\mathrm{pre}}(\cdot|x_{t} ,c),y^{\prime}\sim p^{\diamond}(\cdot|x_{T},c)}[\mathrm{I}(y=y^{\prime})|x_{T} ,c]}_{\mathrm{Additional\,Drift}:=\nabla_{x_{t}}\log p(y|x_{t},c)}\}dt+\sigma( t)dw_{t},$ (4) $$

the marginal distribution of $x_{T}$, i.e., $p(x_{T}|c,y)$, is equal to the target distribution $p_{\gamma=1}(\cdot|c,y)$ (cf. (2)). Here, $\mathbb{P}^{\mathrm{pre}}$ denotes the distribution induced by the pre-trained diffusion model (1).

This lemma suggests that in order to simulate the target distribution (2), we only need to construct SDE (4).
However, applying this method in practice incurs several issues.
First, training the classifier over the time horizon $p(y|x_{t},c)$ necessitates data $\{x^{(i)}_{t},c^{(i)},y^{(i)}\}$ at any time $t\in T$ (or its discretized counterpart). Preparing a large amount of such data with a large pre-trained model becomes cumbersome and can result in suboptimal training. Furthermore, the potential poor performance due to accumulated inaccuracies in drift estimates is a significant concern (Li and van der Schaar, ).

###### Lemma 1 (Doob’s h-transforms (Rogers and Williams, 2000 ) ) .

##### Reconstruction guidance.

To mitigate these issues, several studies (Ho et al., ; Han et al., ; Chung et al., ; Guo et al., )(^3^33We categorize them as reconstruction guidance methods for simplicity. We note that there are many variants.) propose to approximate $p(y|x_{t},c)$ directly via reconstruction, specifically by
$p(y|x_{t},c)=\int p^{\diamond}(y|x_{T},c)p(x_{T}|x_{t},c)dx_{T}\approx p(y|
\hat{x}_{T}(x_{t},c),c)$,
where $\hat{x}_{T}(x_{t},c)$ is the expected denoised sample given $x_{t},c$, i.e., $\hat{x}_{T}(x_{t},c)=\mathbb{E}[x_{T}|x_{t},c]$. Given such an approximation, we only need to learn $p^{\diamond}(y|x_{T},c)$ from data. However, this approximation could be imprecise when $\mathbb{P}(x_{T}|x_{t},c)$ is noisy and difficult to predict reliably (Chung et al., ).

## 4 Conditioning Pre-Trained Diffusion Models with RL

This section provides details on how our method solves the aforementioned goal with methodological motivations.
We begin with a key observation: the conditioning problem can be effectively conceptualized as an RL problem. Building upon this insight, we illustrate our main algorithm.

### 4.1 Conditioning as RL

Recall that our objective is to learn a drift term $g$ in (3) so that the induced marginal distribution at $T$ (i.e., $p^{g}_{T}$) closely matches our target distribution $p_{\gamma}$. To achieve this, we first formulate the problem via the following minimization:

$$ $\displaystyle\mathop{\mathrm{argmin}}_{g}\mathrm{KL}(p^{g}_{T}(\cdot|c,y)\|p_{ \gamma}(\cdot|c,y)).$ $$

With some algebra, we can show that the above optimization problem is equivalent to the following:

$$ $\displaystyle\mathop{\mathrm{argmin}}_{g}\mathbb{E}_{x_{0:T}\sim\mathbb{P}^{g} (\cdot|c,y)}\left[-\gamma\log p^{\diamond}(y|x,c)+\frac{1}{2}\int_{0}^{T}\frac {\|f^{\mathrm{pre}}(s,c,x_{s})-g(s,c,y,x_{s})\|^{2}}{\sigma^{2}(s)}ds\right].$ $$

Here, recall that $\mathbb{P}^{g}$ is the measure induced by SDE (3) with a drift coefficient $g$. Based on this observation, we derive the following theorem.

###### Theorem 1 (Conditioning as RL) .

Consider the following RL problem:

$$ $\displaystyle g^{\star}:=\mathop{\mathrm{argmax}}_{g}\mathbb{E}_{\begin{ subarray}{c}(c,y)\sim\Pi(c,y)\\ x_{0:T}\sim\mathbb{P}^{g}(\cdot|c,y)\end{subarray}}\left[\gamma\log p^{ \diamond}(y|x_{T},c)-\frac{1}{2}\int_{0}^{T}\frac{\|f^{\mathrm{pre}}(s,c,x_{s} )-g(s,c,y,x_{s})\|^{2}}{\sigma^{2}(s)}ds\right],$ (5) $$

where $\Pi\in\Delta(\mathcal{C}\times\mathcal{Y})$.
Significantly,
the marginal distribution $p^{g^{\star}}_{T}$ exactly matches our target distribution:

$$ $\displaystyle\forall(c,y)\in\mathrm{Supp}(\Pi);\,p^{g^{\star}}_{T}(\cdot|c,y)= p^{\gamma}(\cdot|c,y).$ $$

This theorem shows that after solving the RL problem in (5) and obtaining an optimal drift term $g^{\star}$, we are able to sample from the target distribution $p^{\gamma}(\cdot|c,y)$ by following SDE (3) from 0 0 to $T$. In the next section, we explain how to solve (5) in practice.

###### Theorem 1 (Conditioning as RL) .

### 4.2 Algorithm

Inspired by Theorem 1, we provide algorithmic details of our method.
Full pseudocode is listed in Algorithm 1. Our algorithm consists of three major steps, which are introduced below.

Figure: Algorithm 1 Conditioning pre-Trained diffusion models with Reinforcement Learning (CTRL)

##### Step 1: Constructing the augmented model (Line 2).

To add additional conditioning to the pre-trained diffusion model,
it is necessary to enhance the pre-trained model $f^{\mathrm{pre}}(t,c,x;\theta)$. We introduce an augmented model $g(t,c,y,x;\psi)$ with parameters $\psi=[\theta^{\top},\phi^{\top}]^{\top}$. Here, $\psi$ is structured as a combination of the existing parameters $\theta$ and new parameters $\phi$. In fine-tuning, $\psi$ is initialized at $\psi^{\mathrm{ini}}=[{\theta^{\mathrm{pre}}}^{\top},\mathbf{0}^{\top}]$.

Determining the specific architecture of the augmented model involves a tradeoff: adding more new parameters enhances expressiveness but increases training time.
In scenarios where $\mathcal{Y}$ is discrete with cardinality $|\mathcal{Y}|$, the most straightforward solution is to instantiate $\phi$ with a simple linear embedding layer that maps each $y\in\mathcal{Y}$ to its corresponding embedding. These embeddings are then added to every intermediate output in the diffusion SDE (i.e., $x_{t}$ in (3)).
In this way, the original structure is maintained at a maximum level, and direct adoption of all pre-trained weights is guaranteed. Experimentally, we observe that this lightweight modification leads to accurate conditional generations for complex conditioning tasks, as shown in Section 6.

##### Step 2: Training a calibrated classifier with offline data (Line 3).

Using a function class $\mathcal{F}\subset[\mathcal{C}\times\mathcal{X}\to\Delta(\mathcal{Y})]$, such as a neural network, we perform maximum likelihood estimation (MLE):

$$ $\displaystyle\hat{p}(\cdot|x,c):=\mathop{\mathrm{argmax}}_{r\in\mathcal{F}} \sum_{i=1}^{n}\log r(y^{(i)}|x^{(i)},c^{(i)}).$ (6) $$

For example, when $\mathcal{Y}$ is discrete, this loss reduces to the standard cross entropy loss. When $\mathcal{Y}$ is continuous, assuming Gaussian noise, it reduces to a loss in regression problems.

Figure: Algorithm 2 Direct back-propagation for conditioning

##### Step 3: Planning (Line 4).

Equipped with a classifier, we proceed to solve the RL problem (5), which is the key part of the proposed algorithm. As noted
by Black et al. (); Fan et al. (), the diffusion model can be regarded as a special Markov Decision Process (MDP) with known transition dynamics. Thus, it is technically feasible to employ any off-the-shelf RL algorithms.

Inspired by (Clark et al., ; Prabhudesai et al., ), our method, listed in Algorithm 2, is based on direct back-propagation. This method is iterative in nature. During each iteration, we: (1) compute the expectation over trajectories ($\mathbb{E}_{x_{0:T}\sim\mathbb{P}^{g(\cdot;\psi)}}$) using discretization techniques such as Euler-Maruyama; (2) directly optimize the KL-regularized objective function with respect to parameters of the augmented model (i.e., $\psi$).

In practice, such computation might be memory-intensive when there are numerous discretization steps and the diffusion models have a large number of parameters. This is because gradients would need to be back-propagated through the diffusion process. To improve computational efficiency, we recommend employing specific techniques, including (a) only fine-tuning LoRA (Hu et al., ) modules instead of the full diffusion weights, (b) employing gradient checkpointing (Gruslys et al., ; Chen et al., ) to conserve memory, and (c) randomly truncating gradient back-propagation to avoid computing through all diffusion steps (Clark et al., ; Prabhudesai et al., ).

###### Remark 2 (PPO) .

In Algorithm 1, we employ
direct back-propagation (i.e., Algorithm 2) for planning (i.e., solving the RL problem (5)), which necessarily demands the differentiability of the classifier.
If the classifier is non-differentiable, we suggest using Proximal Policy Optimization (PPO) for planning, such as Black et al. (); Fan et al. (). Other parts remain unchanged.

###### Remark 3 (Using classifier-free guidance to adjust guidance strength) .

Throughout the fine-tuning process demonstrated in Algorithm 1 and Algorithm 2, the guidance strength for the additional conditional control (i.e., $y$) is fixed at a specific $\gamma$ (see the target conditional distribution (2)). However, we note that during inference, this guidance strength $\gamma$ can be adjusted—either increased or decreased—using the classifier-free guidance technique.
Details are deferred to  Appendix B.

###### Remark 4 (Choice of exploratory distribution $\Pi$ ) .

According to Theorem 1, it is desired to improve the coverage over $\mathcal{C}\times\mathcal{Y}$ in fine-tuning the model. For example, in practice, if $\mathcal{Y}$ only takes several discrete values, we can sample $y\in\mathcal{Y}$ uniformly from these values as done in Section 6.

###### Remark 2 (PPO) .

###### Remark 3 (Using classifier-free guidance to adjust guidance strength) .

###### Remark 4 (Choice of exploratory distribution Π Π \Pi roman_Π ) .

### 4.3 Source of Errors in CTRL

We elucidate the sources of errors that our algorithm may encounter. This is helpful when comparing our algorithm with an existing algorithm in the next section.

##### Statistical error.

Statistical errors arise during the training of a classifier $\hat{p}(y|x,c)$ from offline data while learning $p^{\diamond}(y|x,c)$. A typical statistical error is given by:

$$ $\displaystyle\mathbb{E}_{(x,c)\sim l^{\mathrm{off}}}[\|\hat{p}(\cdot|x,c)-p^{ \diamond}(\cdot|x,c)\|^{2}_{1}]=O(\mathrm{Cap}(\mathcal{F})/n),$ (7) $$

where $l^{\mathrm{off}}\in\Delta(\mathcal{X}\times\mathcal{C})$ represents the distribution of offline data, and $\mathrm{Cap}(\mathcal{F})$ denotes the size of the function class $\mathcal{F}$ (Wainwright, ).

##### Model-misspecification error.

There are two potential sources. First, when learning $p^{\diamond}(y|x,c)$, we may encounter misspecification errors. Second, during the construction of an augmented model, there could be another misspecification error if the augmented model does not adequately capture the optimal drift $g^{\star}$.

##### Optimization error.

Optimization errors may occur during both the MLE training and the planning process.

## 5 Additional Comparisons with Existing Conditioning Methods

In this section, we further clarify the connections and comparisons between our algorithm and the existing methods.

### 5.1 Comparison to Classifier Guidance

We explore the benefits of CTRL compared to classifier guidance (Dhariwal and Nichol, ). We also provide a theoretical insight that bridges the two approaches.

First, we point out a deep theoretical connection between the two approaches, although they serve very distinct purposes: classifier guidance is primarily an inference technique, whereas our method involves fine-tuning an augmented model. We demonstrate this by deriving the analytical expression for the optimal drift term in the RL problem (5).

###### Lemma 2 (Bridging RL-based conditioning with classifier guidance) .

The optimal drift term $g^{\star}$ for RL problem (5) has the following explicit solution:

$$ $\displaystyle g^{\star}(t,c,y,x_{t})=f^{\mathrm{pre}}(t,c,x_{t})+\sigma^{2}(t) \nabla_{x_{t}}\log{\mathbb{E}_{\mathbb{P}^{\mathrm{pre}}(\cdot|x_{t},c)}\left[ (p^{\diamond}(y|x_{T},c))^{\gamma}|x_{t},c\right]},\quad\forall t\in[0,T]$ $$

This lemma indicates that when $\gamma=1$, the optimal drift $g^{*}$ corresponds to the drift term obtained from Doob’s h-transform (i.e.,  Lemma 1), which is a precise used formula in classifier guidance.

Despite the connection between our algorithm and classifier guidance, which relies on the analytical expression in Lemma 2, they are fundamentally different. Among the issues, the most notable one is that classifier guidance requires us to learn a predictor from $x_{t}$ to $y$ for any $t\in T$, which leads to accumulative inaccuracies from 0 0 to $T$. In contrast, our algorithm directly addresses the RL problem (5) without relying on this analytical expression, thus avoiding the need to learn such predictors over the time horizon.

In  Section 3.2.2, we explore reconstruction guidance-based methods that also aim to circumvent direct predictions from $x_{t}$ to $y$. These methods propose initially mapping $x_{t}$ to a denoised output $x_{T}$ using an approximation $\hat{x}_{T}(x_{t})$, and subsequently relying on $\hat{x}_{T}(x_{t})$ for further computation rather than on each $x_{t}$.
However, this approximation can be imprecise, particularly over a long time horizon. Indeed, as demonstrated in [Theorem 1, (Chung et al., )], inherent errors induced by this approximation persist even without statistical, model-misspecification, or optimization errors, while these approximation errors don’t show up in our algorithm.

###### Lemma 2 (Bridging RL-based conditioning with classifier guidance) .

### 5.2 Comparison to Classifier-Free Guidance

In this subsection, we explain the advancements of CTRL over classifier-free guidance. We begin by detailing how CTRL can harness *conditional independence* to further ease implementation, a capability lacking in classifier-free guidance. Lastly, we discuss the improvements regarding sample (statistical) efficiency.

#### 5.2.1 Leveraging Conditional Independence, Compositionally via CTRL

We detail two scenarios in which our method demonstrates superiority compared to the classifier-free approach by leveraging the conditional independence between inputs and additional controls.

###### Example 1 (Scenario $Y\perp C|X$ ) .

Consider the scenario where a new condition $Y$ is conditionally independent from an existing condition $C$ given $X$, such that $p^{\diamond}(y|x,c)=p(y|x)$. A direct implication is that CTRL can effectively operate using just $(x,y)$ pairs, eliminating the need for $(c,x,y)$ triplets from the offline dataset.

This scenario is frequently observed in practice.
For example, in the case of using the Stable Diffusion pre-trained model (Rombach et al., ) where
$X$ represents an image and
$C$ a text prompt, we may also desire to condition the generations on $Y$: $Y$ being score functions such as compressibility, aesthetic score, or color (Black et al., ).
These scores would depend solely on the image itself while being independent of the prompt. That being said, $Y$ and $C$ are independent conditional on $X$. We further explore this scenario through experimental analysis in  Section 6.1.

###### Example 1 (Scenario Y ⟂ C | X perpendicular-to 𝑌 conditional 𝐶 𝑋 Y\perp C|X italic_Y ⟂ italic_C | italic_X ) .

##### Multi-task conditional generation.

Multi-task conditional generation poses a significant challenge, involving adding multiple controls into pre-trained models. In the following example, we demonstrate how our method can be extended for this multi-task generation.

###### Example 2 (Scenario $Y_{1}\perp Y_{2}|X,C$ ) .

Consider the scenario where two conditions, $Y_{1}$ and $Y_{2}$, exhibit conditional independence given $X$ and $C$, such that $\log p(y_{1},y_{2}|x,c)=\log p(y_{1}|x,c)+\log p(y_{2}|x,c)$. Given the independence, the two classifiers can be effectively trained using $(c,x,y_{1})$ and $(c,x,y_{2})$ triplets respectively. Even more, if $Y_{1}$ and $Y_{2}$ are also independent of $C$ given $X$ (i.e., Example 1), the classifiers can instead utilize $(x,y_{1})$ and $(x,y_{2})$ pairs, significantly simplifying the dataset construction.

This scenario is also prevalently seen in practical applications. Still consider using the Stable Diffusion pre-trained model for an example, where $X$ represents an image and $C$ a text prompt. When conditioning on additional characteristics such as $Y_{1}$ (compressibility) and $Y_{2}$ (color), these attributes depend solely on the image itself, not the associated prompt. Therefore, we may leverage this independence of $Y_{1}$ and $Y_{2}$ from $C$ when conditioned on $X$ to simplify the model’s fine-tuning procedure. The effectiveness of our approach for this setting is further validated experimentally in  Section 6.2.

###### Example 2 (Scenario Y 1 ⟂ Y 2 | X , C perpendicular-to subscript 𝑌 1 conditional subscript 𝑌 2 𝑋 𝐶 Y_{1}\perp Y_{2}|X,C italic_Y start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT ⟂ italic_Y start_POSTSUBSCRIPT 2 end_POSTSUBSCRIPT | italic_X , italic_C ) .

##### Can classifier-free guidance leverage conditional independence?

The applicability of conditional independence in classifier-free guidance, which directly models $p_{\gamma}(\cdot|c,y)$, remains uncertain. For instance, in a scenario where $Y\perp C|X$ as shown in  Example 1, our method requires only pairs $(x,y)$, while classifier-free guidance typically necessitates triplets $(c,x,y)$. Similarly, in cases where $Y_{1}\perp Y_{2}|C,X$ as shown in  Example 2, our approach utilizes triplets $(c,x,y_{1})$ and $(c,x,y_{2})$.
However, as far as we are concerned, quadruples $(c,x,y_{1},y_{2})$ are necessary to employ classifier-free guidance, but obtaining a large volume of such data could become a bottleneck.

#### 5.2.2 Statistical Efficiency

We explain why our approach is more sample-efficient compared to classifier-free guidance. Most importantly, we utilize a pre-trained model to provide access to sampling from $p^{\mathrm{pre}}(x|c)$, which is already trained on large datasets. This allows us to focus solely on modeling the classifier $p^{\diamond}(y|x,c)$ using offline data. Thus, statistical errors from the offline data impact only the classifier learning step (6). In contrast, classifier-free guidance attempts to model the entire $p_{\gamma}(\cdot|c,y)$ directly using offline data. Therefore, it is considered that our method can enhance sample efficiency by learning smaller parts from the offline data.

## 6 Experiments

Figure: (a) Histogram of generations’ labels from $p^{\mathrm{pre}}$.
Refer to caption: x1.png

We compare CTRL with DPS (Chung et al., ), which is a variant of *reconstruction guidance*. We will start by introducing the baseline and describing the experimental settings. For more detailed information on each experiment, such as dataset, architecture, and hyperparameters, please refer to  Appendix D.
Note due to the burden of augmenting data, we do not compare CTRL with the classifier-free guidance method in this work, as detailed in Section 5.2.1.

##### Experimental setup.

In all experiments, we use Stable Diffusion v1.5 (Rombach et al., ) as the pre-trained model $p^{\mathrm{pre}}(x|c)$, here $c$ is a text prompt (e.g., “cat” or “dog”) and $x$ is the corresponding image. Alongside the pre-trained model, we leverage offline datasets comprising $\{x,y\}$ pairs, where $x$ is an image and $y$ denotes its label, such as compressibility and aesthetic score. Since these functions are solely based on image $x$, employing CTRL only demands $\{x,y\}$ pairs instead of triplets $\{c,x,y\}$, thanks to the conditional independence: $Y\perp C|X$ (cf. Example 1).

### 6.1 Conditional Generation on Compressibility

Figure: (a) Training curve
Refer to caption: x2.png

We start by conditioning generations on their file sizes, specifically focusing on compressibility (^4^44It is important to note that, unlike standard tasks in classifier guidance (Chung et al., ), this score is non-differentiable w.r.t. images.). In accordance with Black et al. (), we define compressibility score as the negative file size in kilobytes (kb) of the image after JPEG compression. Keeping the resolution of all generations fixed at $512\times 512$, we ensure that the file sizes are determined solely by the compressibility of the images. Subsequently, denoting compressibility as CP, we define $4$ compressibility labels as follows:
$Y=0$ : $\textbf{CP}<-110.0$;
$Y=1$ : $-110.0\leq\textbf{CP}<-85.0$; $Y=2$ : $-85.0\leq\textbf{CP}<-60.0$; $Y=3$ : $\textbf{CP}\geq-60.0$. Particularly, as depicted in Figure 1(a), generating samples conditioned on $Y=3$ is challenging due to the infrequent occurrence of such samples from the pre-trained model.

##### Results.

Since our conditions are four compressibility levels, we evaluate the performances via (1) obtaining conditional generations by setting $Y\in[0,1,2,3]$ uniformly;
(2) checking if the generations are correct;
(3) computing the mean compressibilities of the generations conditioned on each $Y$.
We plot the mean compressibilities curve in Figure 2(a) and report the evaluation statistics in Table 2(b).

We find that CTRL can precisely generate samples for each targeted condition. In contrast, as illustrated in Figure 1(b), DPS struggles to generate samples accurately conditioned on $Y=1$ and $Y=2$. Therefore, our approach demonstrates a notable improvement over the baseline.
Additional generated images can be found in Appendix D.2.

Figure: (a) Histogram of $1024$ images generated by the pre-trained model.
Refer to caption: x4.png

### 6.2 Multi-task Conditional Generation

We now move on to a more challenging setting: multi-task conditional generation. In this experiment, in addition to compressibility, we *simultaneously* aim to condition the generations on their aesthetic pleasingness. Following prior research (Black et al., ; Fan et al., ; Uehara et al., ), we employ an aesthetic scorer implemented as a linear MLP on top of the CLIP embeddings (Radford et al., ), which is trained on more than $400$k human evaluations.

In this experiment, by leveraging conditional independence of $Y_{1}$ and $Y_{1}$ given $X$ (cf.  Example 2), we aim to fine-tune the diffusion model to generate samples with *compositional* conditions. Specifically, denoting compressibility as CP and aesthetic score as AS, we define four compositional conditions as follows: $Y=0$ : $\textbf{AS}<5.7,~{}\textbf{CP}<-70$; $Y=1$ : $\textbf{AS}<5.7,~{}\textbf{CP}\geq-70$; $Y=2$ : $\textbf{AS}\geq 5.7,~{}\textbf{CP}<-70$; $Y=3$ : $\textbf{AS}\geq 5.7,~{}\textbf{CP}\geq-70$. Particularly, as depicted in Figure 3(a), generating samples conditioned on $Y=1$ or $Y=3$ is challenging due to the infrequent occurrence of such samples from the pre-trained model.

##### Results.

We follow the same evaluation procedure as described in Section 6.1, and present evaluation statistics in Table 3(b). Again, we observe that CTRL can accurately generate samples for each specified condition.
Our approach exhibits a significant improvement over the baseline. More visualizations are provided in  Appendix D.2.

## 7 Summary

We introduce an RL-based fine-tuning approach for conditioning pre-trained diffusion models on new additional labels. Compared to classifier-free guidance, our proposed method uses the offline dataset more efficiently and allows for leveraging the conditional independence assumption, thereby greatly simplifying the construction of the offline dataset.

While current experiments are restricted to the computer vision domain where sample efficiency may not be problematic, our next step involves applying our algorithm to scientific domains.

## Appendix A Training Diffusion Models

In standard diffusion models, given a training dataset $\{x^{\langle j\rangle}\}\sim p_{\mathrm{data}}(\cdot)$, the goal is to construct a transport that maps noise distribution and data distribution $p_{\mathrm{data}}\in\Delta(\mathcal{X})$ ($\mathcal{X}=\ \mathbb{R}^{d}$). More specifically, suppose that we have an SDE (^5^55In standard diffuson models, the direction is reversed, i.e., $x_{T}$ corresponds to the noise distribution.):

$$ $\displaystyle dx_{t}=f(t,x_{t};\theta)dt+\sigma(t)dw_{t},$ (8) $$

where $f:[0,T]\times\mathbb{R}^{d}\to\mathbb{R}^{d}$ is a drift coefficient, $\sigma:[0,T]\to\mathbb{R}$ is a diffusion coefficient, $w_{t}$ is $d$-dimensional Brownian motion, and initial state $x_{0}\sim p_{\mathrm{ini}}$ where $p_{\mathrm{ini}}\in\Delta(\mathcal{X})$ denotes the initial distribution. By denoting the marginal distribution at time $T$ by $p^{\theta}_{T}(x)$, a standard goal in training diffusion models is to learn the parameter $\theta$ so that $p^{\theta}_{T}(x)\approx p_{\mathrm{data}}$. This means we can (approximately) sample from $p_{\mathrm{data}}$ by following the SDE (8) from 0 0 to $T$.

To train diffusion models, we first introduce a (fixed) *forward* reference SDE, which gradually adds noise to $p_{\mathrm{data}}$:

$$ $\displaystyle dz_{t}=\bar{f}(t,z_{t})dt+\bar{\sigma}(t)dw_{t},\quad z_{0}\sim p _{\mathrm{data}},$ (9) $$

where $\bar{f}:[0,T]\times\mathbb{R}^{d}\to\mathbb{R}^{d}$ is a drift coefficient, $\bar{\sigma}:[0,T]\to\mathbb{R}$ is a diffusion coefficient. An example is the classical denoising diffusion model (Ho et al., ), also known as the variance-preserving (VP) process, which sets $\bar{f}=-0.5z_{t},\bar{\sigma}=1$.

Now, we consider the time-reversal SDE (Anderson, ), which reverses the direction of SDE while keeping the marginal distribution, as follows:

$$ $\displaystyle dx_{t}=\left\{-\bar{f}(T-t,x_{t})+\nabla\log q_{T-t}(x_{t}) \right\}dt+\bar{\sigma}(T-t)dw_{t},\quad x_{0}\sim\mathcal{N}(0,I_{d}).$ (10) $$

Here, $q_{t}(\cdot)$ denotes the marginal distribution at time $t$ for the distribution induced by the reference SDE, and $\nabla\log q_{T-t}(x_{t})$ means a derivative w.r.t. $x_{t}$, which is often referred to as the *score function*. Furthremore, when the time horizon $T$ is sufficiently large, $z_{t}$ follows Gaussian noise distribution $\mathcal{N}(0,I_{d})$.
Hence, if we could learn the score function, by following the SDE (10) starting from Gaussian noise, we can sample from the data distribution.

Then, we aim to learn the score function from the data. By comparing the time-reversal SDE with the original SDE, a natural parameterization is:

$$ $f(t,x_{t};\theta)=-\bar{f}(T-t,x_{t})+s(T-t,x_{t};\theta),\quad\sigma(t)=\bar{ \sigma}(T-t),$ $$

where $s(T-t,x_{t};\theta)$ is the parametrized neural network introduced to approximate the score function $\nabla\log q_{T-t}(x_{t})$.
Here, we can leverage the analytical form of the conditional distribution $q_{T-t|0}(\cdot|\cdot)$ (which is a Gaussian distribution derived from the reference SDE). This approach enables us to tackle the approximation problem via regression:

$$ $\displaystyle\hat{\theta}=\mathop{\mathrm{argmin}}_{\theta}\mathbb{E}_{t\in[0, T],z_{0}\sim p_{\text{data}},z_{t}\sim q_{t|0}(z_{0})}\left[\lambda(t)\left\|s (t,z_{t};\theta)-\nabla_{z_{t}}\log q_{t|0}(z_{t}|z_{0})\right\|^{2}\right],$ (11) $$

where $\lambda:[0,T]\to\mathbb{R}$ is a weighting function.

## Appendix B Inference Technique in Classfier-free Guidance

Although the fine-tuning process sets the guidance level for the additional conditioning (i.e., $y$) at a specific $\gamma$, classifier-free guidance makes it possible to adjust the guidance strength freely during inference. Recall that the augmented model is constructed as:
$g(t,c,y,x;\psi)$ where $\psi=[\theta^{\top},\phi^{\top}]^{\top}$. Suppose we have obtained a drift term $\hat{g}$, parametrized by $\hat{\psi}=[\hat{\theta}^{\top},\hat{\phi}^{\top}]^{\top}$ from running Algorithm 1. In inference, we may alter the guidance levels by using the following drift term in the SDE (3)

$$ $\displaystyle\quad g_{\gamma_{1},\gamma_{2}}(t,c,y,x_{t})$ $\displaystyle=\underbrace{g(t,\emptyset,\emptyset,x_{t};\hat{\psi})+\gamma_{1} (g(t,c,\emptyset,x_{t};\hat{\psi})-g(t,\emptyset,\emptyset,x_{t};\hat{\psi}))} _{\text{Term $1$: pre-trained diffusion model conditioned on $\mathcal{C}$}}+ \underbrace{\gamma_{2}(g(t,c,y,x_{t};\hat{\psi})-g(t,c,\emptyset,x_{t};\hat{ \psi}))}_{\text{Term $2$: additional conditioning on $\mathcal{Y}$}}$ $$

where $\emptyset$ indicates the unconditional on $\mathcal{Y}$ or on $\mathcal{C}$. In the above, both $\gamma_{1}$ and $\gamma_{2}$ do not necessarily need to equal $\gamma$. They can be adjusted respectively to reflect guidance strength levels for two conditions.

## Appendix C Proofs

### C.1 Important Lemmas

We first introduce several important lemmas to prove our main statement.

First, recall that $\mathbb{P}^{g}(\cdot|c,y)$ is the induced distribution by the SDE:

$$ $\displaystyle dx_{t}=g(t,c,y,x_{t})dt+\sigma(t)dw_{t},\quad x_{0}=x_{\mathrm{ ini}}$ $$

over $\mathcal{K}$ conditioning on $c$ and $y$. Similarly, denote $\mathbb{P}^{\mathrm{pre}}(\cdot|c)$ by the induced distribution by the SDE:

$$ $\displaystyle dx_{t}=f^{\mathrm{pre}}(t,c,x_{t})dt+\sigma(t)dw_{t},\quad x_{0} =x_{\mathrm{ini}}$ $$

over $\mathcal{K}$ conditioning on $c$.

###### Lemma 3 (KL-constrained reward) .

The objective function in (5) is equivalent to

$$ $\displaystyle\mathrm{obj}=\mathbb{E}_{(c,y)\sim\Pi,\mathbb{P}^{g}(\cdot|c,y)}[ \gamma\log{p^{\diamond}(y|x_{T},c)}-\mathrm{KL}(\mathbb{P}^{g}(\cdot|c,y)\| \mathbb{P}^{\mathrm{pre}}(\cdot|c)].$ (12) $$

###### Proof.

We calculate the KL divergence of $\mathbb{P}^{g}$ and $\mathbb{P}^{\mathrm{pre}}$ as below

$$ $\displaystyle\mathrm{KL}(\mathbb{P}^{g}(\cdot|c,y)\|\mathbb{P}^{\mathrm{pre}}( \cdot|c)=\mathbb{E}_{x_{0:T}\sim\mathbb{P}^{g}(\cdot|c,y)}\left[\int_{0}^{T} \frac{1}{2}\frac{\|g(t,c,y,x_{t})-f^{\mathrm{pre}}(t,c,x_{t})\|^{2}}{\sigma^{2 }(t)}dt\right].$ (13) $$

This is because

$$ $\displaystyle\quad\mathrm{KL}(\mathbb{P}^{g}(\cdot|c,y)\|\mathbb{P}^{\mathrm{ pre}}(\cdot|c))$ $\displaystyle=\mathbb{E}_{\mathbb{P}^{g}(\cdot|c,y)}\left[\frac{d\mathbb{P}^{g }(\cdot|c,y)}{d\mathbb{P}^{\mathrm{pre}}(\cdot|c)}\right]$ $\displaystyle=\mathbb{E}_{\mathbb{P}^{g}(\cdot|c,y)}\left[\int_{0}^{T}\frac{1} {2}\frac{\|g(t,c,y,x_{t})-f^{\mathrm{pre}}(t,c,x_{t})\|^{2}}{\sigma^{2}(t)}dt+ \int_{0}^{T}\{g(t,c,y,x_{t})-f^{\mathrm{pre}}(t,c,x_{t})\}dw_{t}\right]$ (Girsanov theorem) $\displaystyle=\mathbb{E}_{\mathbb{P}^{g}(\cdot|c,y)}\left[\int_{0}^{T}\frac{1} {2}\frac{\|g(t,c,y,x_{t})-f^{\mathrm{pre}}(t,c,x_{t})\|^{2}}{\sigma^{2}(t)}dt \right].$ (Martingale property of Itô integral) $$

Therefore, the objective function in (5) is equivalent to

$$ $\displaystyle\mathrm{obj}=\mathbb{E}_{(c,y)\sim\Pi,\mathbb{P}^{g}(\cdot|c,y)}[ \gamma\log{p^{\diamond}(y|x_{T},c)}-\mathrm{KL}(\mathbb{P}^{g}\|\mathbb{P}^{ \mathrm{pre}})].$ (14) $$

∎

###### Lemma 3 (KL-constrained reward) .

###### Proof.

##### Optimal value function.

For the RL problem (5), it is beneficial to introduce the optimal optimal value function $v^{\star}_{t}(x|c,y)$ at any time $t\in[0,T]$, given $x_{t}=x$, conditioned on parameters $c$ and $y$ defined as:

$$ $\displaystyle v^{\star}_{t}(x|c,y)=\max_{g}\mathbb{E}\left[\gamma\log p^{ \diamond}(y|x_{T},c)-\frac{1}{2}\int_{t}^{T}\frac{\|f^{\mathrm{pre}}(s,c,x_{s} )-g(s,c,y,x_{s})\|^{2}}{\sigma^{2}(s)}\,ds\;\Bigg{|}\;x_{t}=x,c,y\right].$ (15) $$

Specifically, we note that $v_{T}(x|c,y)=\gamma\log{p^{\diamond}(y|x,c)}$ represents the terminal reward function (i.e., a loglikelihood in our MDP), while $v_{0}$ represents the original objective function (5) that integrates the entire trajectory’s
KL divergence along with the terminal reward.

Below we derive the optimal value function in analytical form.

###### Lemma 4 (Feynman–Kac Formulation) .

At any time $t\in[0,T]$, given $x_{t}=x$, and conditioned on $c$ and $y$, we have the optimal value function $v_{t}^{*}(x|c,y)$ (induced by the optimal drift term $g^{*}$) as follows

$$ $\exp\left(v^{\star}_{t}(x|c,y)\right)=\mathbb{E}_{\mathbb{P}^{\mathrm{pre}}( \cdot|c)}\left[(p^{\diamond}(y|x_{T},c))^{\gamma}|x_{t}=x,c\right].$ $$

###### Proof.

From the Hamilton–Jacobi–Bellman (HJB) equation, we have

$$ $\displaystyle\max_{u}\left\{\frac{\sigma^{2}(t)}{2}\sum_{i}\frac{d^{2}v^{\star }_{t}(x|c,y)}{dx^{[i]}dx^{[i]}}+g\cdot\nabla v^{\star}_{t}(x|c,y)+\frac{dv^{ \star}_{t}(x|c,y)}{dt}-\frac{\|g-f^{\mathrm{pre}}\|^{2}_{2}}{2\sigma^{2}(t)} \right\}=0.$ (16) $$

where $x^{[i]}$ is a $i$-th element in $x$. Hence, by simple algebra, we can prove that the optimal drift term satisfies

$$ $\displaystyle g^{\star}(t,c,y,x)=f^{\mathrm{pre}}(t,c,x)+\sigma^{2}(t)\nabla v ^{\star}_{t}(x|c,y).$ $$

By plugging the above into the HJB equation (16), we get

$$ $\displaystyle\frac{\sigma^{2}(t)}{2}\sum_{i}\frac{d^{2}v^{\star}_{t}(x|c,y)}{ dx^{[i]}dx^{[i]}}+f^{\mathrm{pre}}\cdot\nabla v^{\star}_{t}(x|c,y)+\frac{dv^{ \star}_{t}(x|c,y)}{dt}+\frac{\sigma^{2}(t)\|\nabla v^{\star}_{t}(x|c,y)\|^{2}_ {2}}{2}=0,$ (17) $$

which characterizes the optimal value function. Now, using (17), we can show

$$ $\displaystyle\frac{\sigma^{2}(t)}{2}\sum_{i}\frac{d^{2}\exp(v^{\star}_{t}(x|c, y))}{dx^{[i]}dx^{[i]}}+f^{\mathrm{pre}}\cdot\nabla\exp(v^{\star}_{t}(x|c,y))+ \frac{d\exp(v^{\star}_{t}(x|c,y))}{dt}$ $\displaystyle=\exp\left(v^{\star}_{t}(x|c,y)\right)\times\left\{\frac{\sigma^{ 2}(t)}{2}\sum_{i}\frac{d^{2}v^{\star}_{t}(x|c,y)}{dx^{[i]}dx^{[i]}}+f^{\mathrm {pre}}\cdot\nabla v^{\star}_{t}(x|c,y)+\frac{dv^{\star}_{t}(x|c,y)}{dt}+\frac{ \sigma^{2}(t)\|\nabla v^{\star}_{t}(x|c,y)\|^{2}_{2}}{2}\right\}$ $\displaystyle=0.$ $$

Therefore, to summarize, we have

$$ $\displaystyle\frac{\sigma^{2}(t)}{2}\sum_{i}\frac{d^{2}\exp(v^{\star}_{t}(x|c, y))}{dx^{[i]}dx^{[i]}}+f^{\mathrm{pre}}\cdot\nabla\exp(v^{\star}_{t}(x|c,y))+ \frac{d\exp(v^{\star}_{t}(x|c,y))}{dt}=0,$ (18) $\displaystyle v^{\star}_{T}(x|c,y)=\gamma\log{p^{\diamond}(y|x,c)}.$ (19) $$

Finally, by invoking the Feynman-Kac formula (Shreve et al., ), we obtain the conclusion:

$$ $\displaystyle\exp\left(v^{\star}_{t}(x|c,y)\right)=\mathbb{E}_{\mathbb{P}^{ \mathrm{pre}}(\cdot|x_{t},c)}\left[(p^{\diamond}(y|x_{T},c))^{\gamma}|x_{t}=x, c\right].$ $$

∎

###### Lemma 4 (Feynman–Kac Formulation) .

###### Proof.

### C.2 Proof of Theorem 1

Firstly, we aim to show that the optimal conditional distribution over $\mathcal{K}$ on $c$ and $y$ (i.e., $\mathbb{P}^{g^{\star}}(\tau|c,y)$) is equivalent to

$$ $\displaystyle\frac{\mathbb{P}^{\mathrm{pre}}(\tau|c)(p^{\diamond}(y|x_{T},c))^ {\gamma}}{C(c,y)},\quad C(c,y):=\exp(v^{\star}_{0}(x_{0}|c,y))).$ $$

To do that, we need to check that the above is a valid distribution first. This is indeed valid because the above is decomposed into

$$ $\displaystyle\underbrace{\frac{(p^{\diamond}(y|x_{T},c))^{\gamma}\cdot\mathbb{ P}^{\mathrm{pre}}(x_{T}|c)}{C(c,y)}}_{(\alpha 1)}\times\underbrace{\mathbb{P}^ {\mathrm{pre}}(\tau|c,x_{T})}_{(\alpha 2)},$ (20) $$

and both $(\alpha 1),(\alpha 2)$ are valid distributions. Especially, for the term $(\alpha 1)$, we observe

$$ $\displaystyle C(c,y)=\int(p^{\diamond}(y|x_{T},c))^{\gamma}d\mathbb{P}^{ \mathrm{pre}}(x_{T}|c))=\mathbb{E}_{\mathbb{P}^{\mathrm{pre}}(\cdot|c)}[(p^{ \diamond}(y|x_{T},c))^{\gamma}]=\exp(v^{\star}_{0}(x_{0}|c,y)).$ (cf. Lemma 4 ) $$

Now, after checking (20) is a valid distribution, we calculate the KL divergence:

$$ $\displaystyle\mathrm{KL}\left(\mathbb{P}^{g^{\star}}(\tau|c,y)\bigg{\|}\frac{ \mathbb{P}^{\mathrm{pre}}(\tau|c)(p^{\diamond}(y|x_{T},c))^{\gamma}}{C(c,y)}\right)$ $\displaystyle=\mathrm{KL}(\mathbb{P}^{g^{\star}}(\tau|c,y)\|\mathbb{P}^{ \mathrm{pre}}(\tau|c))-\mathbb{E}_{\mathbb{P}^{g^{\star}}(\cdot|c,y)}\left[ \gamma\log{p^{\diamond}(y|x_{T},c)}-\log C(c,y)\right]$ $\displaystyle=\mathbb{E}_{\mathbb{P}^{g^{\star}}(\cdot|c,y)}\left[\left\{\int_ {0}^{T}\frac{1}{2}\frac{\|g^{\star}(t,c,y,x_{t})-f^{\mathrm{pre}}(t,c,x_{t})\| ^{2}}{\sigma^{2}(t)}\right\}dt-\gamma\log{p^{\diamond}(y|x_{T},c)}+\log C(c,y)\right]$ (cf. KL divergence ( 13 )) $\displaystyle=-v^{\star}_{0}(x_{0}|c,y)+\log C(c,y).$ (Definition of optimal value function) $$

Therefore,

$$ $\displaystyle\mathrm{KL}\left(\mathbb{P}^{g^{\star}}(\tau|c,y)\Big{\|}\frac{ \mathbb{P}^{\mathrm{pre}}(\tau|c)(p^{\diamond}(y|x_{T},c))^{\gamma}}{C(c,y)} \right)=-v^{\star}_{0}(x_{0}|c,y)+\log C(c,y)=0.$ $$

Hence,

$$ $\displaystyle\mathbb{P}^{g^{\star}}(\tau|c,y)=\frac{\mathbb{P}^{\mathrm{pre}}( \tau|c)(p^{\diamond}(y|x_{T},c))^{\gamma}}{C(c,y)}.$ $$

##### Marginal distribution at t 𝑡 t italic_t .

Finally, consider the marginal distribution at $t$. By marginalizing before $t$, we get

$$ $\displaystyle\mathbb{P}^{\mathrm{pre}}(\tau_{[t,T]}|c)\times(p^{\diamond}(y|x_ {T},c))^{\gamma}/C(c,y).$ $$

Next, by marginalizing after $t$,

$$ $\displaystyle\mathbb{P}^{\mathrm{pre}}_{t}(x|c)/C(c,y)\times\mathbb{E}_{ \mathbb{P}^{\mathrm{pre}}(\cdot|c)}[(p^{\diamond}(y|x_{T},c))^{\gamma}|x_{t}=x ,c].$ $$

Using Feynman–Kac formulation in Lemma 4, this is equivalent to

$$ $\displaystyle\mathbb{P}^{\mathrm{pre}}_{t}(x|c)\exp(v^{\star}_{t}(x|c,y))/C(c, y).$ $$

##### Marginal distribution at T 𝑇 T italic_T .

We marginalize before $T$. We have the following

$$ $\displaystyle\mathbb{P}^{\mathrm{pre}}_{T}(x|c)(p^{\diamond}(y|x_{T},c))^{ \gamma}/C(c,y).$ $$

### C.3 Proof of Lemma 2

Recall $g^{\star}(t,c,y,x)=f^{\mathrm{pre}}(t,c,x)+\sigma^{2}(t)\times\nabla_{x}v^{
\star}_{t}(x|c,y)$ from the proof of Lemma 4, we have

$$ $\displaystyle g^{\star}(t,c,y,x)=f^{\mathrm{pre}}(t,c,x)+\sigma^{2}(t)\times \nabla_{x}\log{\mathbb{E}_{\mathbb{P}^{\mathrm{pre}}(\cdot|c)}\left[(p^{ \diamond}(y|x_{T},c))^{\gamma}|x_{t}=x,c\right]}.$ $$

## Appendix D Experiment Details

Below, we explain the training details and list hyperparameters in Table 4.

### D.1 Implementation Details

We use 4 A100 GPUs for all the image tasks.
We use the AdamW optimizer (Loshchilov and Hutter, ) with $\beta_{1}=0.9,\beta_{2}=0.999$ and weight decay of $0.1$.
To ensure consistency with previous research, in fine-tuning, we also employ training prompts that are uniformly sampled from $50$ common animals (Black et al., ; Prabhudesai et al., ).

##### Construction of the augmented score model.

An important engineering aspect is how to craft the augmented score model architecture. For most of the diffusion models, the most natural and direct technique of adding another conditioning control is (1) augmenting the score prediction networks by incorporating additional linear embeddings, while using the existing neural network architecture and weights for all other parts. In our setting, we introduce a linear embedding layer that maps $|\mathcal{Y}|+1$ class labels to embeddings in $\mathbb{R}^{d}$, where $d$ is the same dimension as intermediate diffusion states. Among all embeddings, the first $|\mathcal{Y}|$ embeddings correspond to $|\mathcal{Y}|$ conditions of our interest, whereas the last one represents the unconditional category (i.e., NULL conditioning) (2) for any $y\in\mathcal{Y}$, the corresponding embedding is added to the predicted score in the forward pass. During fine-tuning, the embeddings are initialized as zeros. We only fine-tune the first $|\mathcal{Y}|$ embeddings, and freeze the last one at zero as it is the unconditional label.

We note that, while it is possible to add additional conditioning by reconstructing the score networks like ControlNet (Zhang et al., ), in practice it is often desired to make minimal changes to the architecture of large diffusion models, e.g., Stable Diffusion (Rombach et al., ) to avoid the burdensome re-training. It is especially important to leverage pre-trained diffusion models in our setting where the offline dataset is limited, therefore a total retraining of model parameters can be struggling.

##### Sampling.

We use the DDIM sampler with $50$ diffusion steps (Song et al., ).
Since we need to back-propagate the gradient of rewards through both the sampling process producing the latent representation and the VAE decoder
used to obtain the image, memory becomes a bottleneck. We employ two designs to alleviate memory usage following  Clark et al. (); Prabhudesai et al. ():
(1) Fine-tuning low-rank adapter (LoRA) modules (Hu et al., ) instead of tuning the original diffusion
weights, and (2) Gradient checkpointing for computing partial derivatives on demand (Gruslys et al., ; Chen et al., ). The two designs make it possible to back-propagate gradients through all 50 diffusing steps in terms of hardware.

**Table 2: Architecture of compressibility classifier**
|  | Input Dimension | Output Dimension | Explanation |
| --- | --- | --- | --- |
| 1 | $C\times H\times W$ | $64\times H\times W$ | ResidualBlock (Conv2d(3, 64, 3x3), BN, ReLU) |
| 2 | $64\times H\times W$ | $128\times\frac{H}{2}\times\frac{W}{2}$ | ResidualBlock (Conv2d(64, 128, 3x3), BN, ReLU) |
| 3 | $128\times\frac{H}{2}\times\frac{W}{2}$ | $256\times\frac{H}{4}\times\frac{W}{4}$ | ResidualBlock (Conv2d(128, 256, 3x3), BN, ReLU) |
| 4 | $256\times\frac{H}{4}\times\frac{W}{4}$ | $256\times 1\times 1$ | AdaptiveAvgPool2d (1, 1) |
| 5 | $256\times 1\times 1$ | $256$ | Flatten |
| 6 | $256$ | num_classes | Linear |

**Table 3: Architecture of aesthetic score classifier**
| # | Layer Type | Input Dimension | Output Dimension |
| --- | --- | --- | --- |
| 1 | Linear | $768$ | $1024$ |
| 2 | Dropout | - | - |
| 3 | Linear | $1024$ | $128$ |
| 4 | Dropout | - | - |
| 5 | Linear | $128$ | $64$ |
| 6 | Dropout | - | - |
| 7 | Linear | $64$ | $16$ |
| 8 | Linear | $16$ | num_classes |

**Table 4: Training hyperparameters.**
| Hyperparameter | compressibility (Section 6.1) | multi-task (Section 6.2) |
| --- | --- | --- |
| Classifier-free guidance weight on prompts (i.e., $c$) | $7.5$ | $7.5$ |
| $\gamma$ (i.e., strength of the additional guidance on $y$) | $10$ | $10$ |
| DDIM steps | $50$ | $50$ |
| Truncated back-propagation step | $K\sim\text{Uniform}(0,50)$ | $K\sim\text{Uniform}(0,50)$ |
| Learning rate for LoRA modules | $1e^{-3}$ | $3e^{-4}$ |
| Learning rate for the linear embeddings | $1e^{-2}$ | $1e^{-2}$ |
| Batch size (per gradient update) | $256$ | $512$ |
| Number of gradient updates per epoch | $2$ | $2$ |
| Epochs | $15$ | $60$ |

##### Training classifiers.

In our experiments, we leverage conditional independence for both compressibility and aesthetic scores tasks. Therefore, we only demand data samples $\{x_{i},y_{i}\})$ in order to approximate conditional classifier $p(y|x,c)$. Specifically,

- •
compressibility: the classifier is implemented as a 3-layer convolutional neural network (CNN) with residual connections and batch normalizations on top of
the raw image space. The offline dataset is constructed by labeling a subset of $10$k images of the AVA dataset (Murray et al., ), employing JPEG compression. We train the network using Adam optimizer for $100$ epochs. Detailed architecture of the oracle can be found in Table 2.
- •
aesthetic scores: the classifier is implemented as an MLP on top of CLIP embeddings (Radford et al., ). To train the classifier, we use the full AVA dataset (Murray et al., ) which includes more than 250k human evaluations. The specific neural network instruction is listed in Table 3.

Note that in training both classifiers, we split the dataset with $80\%$ for training and $20\%$ for validation. After training, we use the validation set to perform temperature scaling calibration Guo et al. ().

### D.2 Additional Results

##### More generated images.

We provide more generated samples to illustrate the performances in  Figure 4 and  Figure 5.

Figure: Figure 4: More images generated by CTRL in the compressibility task.
Refer to caption: x6.png

Figure: Figure 5: More images generated by CTRL in the multi-task conditional generation.
Refer to caption: x7.png

**Table 5: Results of DPS conditioned on compressibility. Essentially, guidance level$=0$ indicates that the generations are unconditional on $Y$.**
| Conditional control ($Y$) | Guidance level ($\gamma$) | Accuracy $\uparrow$ | Mean score |
| --- | --- | --- | --- |
| 0 0 | 0 0 | $0.43$ | $-110$ |
| 0 0 | $5$ | $0.64$ | $-157.4$ |
| 0 0 | $7.5$ | $0.62$ | $-148.2$ |
| 0 0 | $10$ | $0.52$ | $-156.5$ |
| 0 0 | $20$ | $0.66$ | $-152.5$ |
| 0 0 | $50$ | $0.69$ | $-151.6$ |
| $1$ | 0 0 | $0.45$ | $-110$ |
| $1$ | $5$ | $0.14$ | $-152.0$ |
| $1$ | $7.5$ | $0.12$ | $-189.7$ |
| $1$ | $10$ | $0.08$ | $-163.1$ |
| $1$ | $20$ | $0.06$ | $-169.5$ |
| $1$ | $50$ | 0 0 | $-194.0$ |
| $2$ | 0 0 | $0.13$ | $-110$ |
| $2$ | $5$ | $0.02$ | $-104.9$ |
| $2$ | $7.5$ | $0.10$ | $-122.1$ |
| $2$ | $10$ | $0.12$ | $-111.3$ |
| $2$ | $20$ | $0.08$ | $-157.6$ |
| $2$ | $50$ | $0.08$ | $-173.5$ |
| $3$ | 0 0 | 0 0 | $-110$ |
| $3$ | $5$ | $0.46$ | $-71.7$ |
| $3$ | $7.5$ | $0.46$ | $-67.6$ |
| $3$ | $10$ | $0.53$ | $-65.6$ |
| $3$ | $20$ | $0.26$ | $-112.4$ |
| $3$ | $50$ | $0.32$ | $-121.5$ |

##### More Result for DPS (Chung et al., 2022 ) .

For completeness, we provide more experimental results and analyses for the DPS baseline conditioned on compressibility.

In Table 5, we present the classification statistics for generations across different conditions ($y$) and guidance levels ($\gamma$). Recall that the four conditions are defined as follows:
$Y=0$ :  $\textbf{CP}<-110.0$, $Y=1$ :  $-110.0\leq\textbf{CP}<-85.0$, $Y=2$ :  $-85.0\leq\textbf{CP}<-60.0$, $Y=3$ :  $\textbf{CP}\geq-60.0$.

Our analysis reveals several key insights:

- 1.
For $Y=0$ and $Y=3$, the accuracy of the generations improves as the guidance signal strength increases. This indicates a clear positive correlation between the guidance level and the accuracy of generation.
- 2.
Conversely, for intermediate $Y=1$ and $Y=2$, guidance signals decrease generation accuracy compared to the pre-trained model, suggesting difficulty in maintaining accuracy within these specific compressibility intervals. The challenge in generating samples with medium compressibility scores lies in hand-picking the guidance strength. For instance, generating samples conditioned on $Y=2$ requires compressibility scores between $-85$ and $-60$, making it difficult to apply optimal guidance without overshooting or undershooting the target values.
- 3.
Regarding the mean scores, distinct patterns are observed across different conditions and guidance levels:
•
For $Y=0$, mean scores become more negative with increasing guidance levels.
•
For $Y=1$, mean scores consistently drop with increasing guidance levels.
•
For $Y=2$, mean scores initially improve slightly with increasing guidance levels but show a marked decline at $\gamma=20$ and $\gamma=50$, indicating a challenge in achieving the desired compressibility range.
•
For $Y=3$, mean scores improve significantly with increased guidance, showing the best results at $\gamma=10$, but then become more negative at higher guidance levels.

In summary, these observations suggest that while guidance can be beneficial for improving accuracy in extreme compressibility levels ($Y=0$ and $Y=3$), this method struggles with intermediate conditions ($Y=1$ and $Y=2$) due to the narrow range of acceptable scores and the non-linear effects of guidance strength on generation quality.

For each conditional control, samples are generated by choosing the best $\gamma$ according to Table 5. We report the evaluation statistics in Table 2(b), and provide the confusion matrix in  Figure 1(b)

## Peer Reviews

confidence: 2
contribution: 2
initial_rating: 5
presentation: 3
questions: + While there are indeed numerous works that frame the reverse diffusion process as an MDP and use reinforcement learning to optimize the diffusion model, I am concerned that such optimization may lack sufficient constraints. To clarify, consider two consecutive time steps, $t$ and $t+1$ in the forward diffusion process, we have the following equation:
 $$p_{t+1}(x_{t+1})=\int p_t(x_t)\mathcal{N}(x_{t+1}; \sqrt{1-\beta_{t+1}}, \beta_{t+1}I)$$
Now, suppose we apply CTRL to fine-tune the reverse process. Since neural networks are used to approximate the score functions and are updated directly to maximize the objective without explicit constraints, it is likely that the above relationship will no longer hold after optimization. In other words, following reinforcement learning, the forward process might no longer correspond to the original stochastic differential equation (SDE). Could the authors provide insights into this concern? Specifically, will such a deviation impact the performance or generalization capability of the diffusion model?
soundness: 3
strengths: + This paper is well-organized and the presentation is clear. 

+ Unlike classifier guidance and classifier-free guidance, CTRL approaches controllable generation from a distinct perspective by framing conditional generation as a reinforcement learning problem and fine-tuning the diffusion model over the reverse diffusion process. While a few prior works (e.g., [1]) have explored similar perspectives, this work is distinguished by its use of a KL-regularized RL framework, deriving the analytical form of KL regularization through Girsanov’s theorem. Additionally, I especially appreciate that the authors included a discussion in the main text on the relationship between CTRL, classifier guidance, and classifier-free guidance, providing valuable context for understanding the nuances of these methods.

[1] Kevin Black, Michael Janner, Yilun Du, Ilya Kostrikov, Sergey Levine. Training Diffusion Models with Reinforcement Learning.
summary: This paper introduces an RL-based approach for integrating conditional control into pre-trained diffusion models. The primary contribution of the proposed algorithm is its RL objective, which is coupled with an analytical form of the KL divergence, allowing for seamless optimization using standard reinforcement learning algorithms. The authors benchmarked the proposed method against established controllable generation techniques, specifically classifier guidance and classifier-free guidance, and empirically demonstrated the effectiveness of the CTRL.
weaknesses: + Theoretically, an RL formulation allows us to view equation (5) as a sequential decision-making problem, enabling the learning of value functions for each diffusion time step $s$ and policy optimization based on these estimated values. However, in this paper (Algorithm 1), the authors instead use the diffusion model to roll out a diffusion path, preserving the gradient of each intermediate $x_t$ (assuming this interpretation is correct), computing the loss, and finally backpropagating the gradient across the entire diffusion path. This approach likely demands higher GPU memory, as it requires gradient preservation throughout the generation process, as well as increased computation, given the need to backpropagate gradients through the model multiple times. These requirements could limit the scalability and broader applicability of the proposed method.

+ Besides, in order to achieve controllable generation, classfier guidance and classifier-free guidance need to train either the classifier or a finetuned model, while CTRL needs to train both of them. This complicates the overall procedure.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
questions: See weaknesses.
soundness: 3
strengths: 1. The problem the authors attempt to address—adding conditional control to pre-trained diffusion models—is highly valuable. The author’s use of RL modeling for this problem is novel, and the derivation process is both reasonable and rigorous.
2. Compared to previous methods, CTRL has more flexible requirements for offline datasets and does not require training a predictor $x_t\rightarrow y$, resulting in a broader range of applications and a more controllable generation process.
3. In the experimental section, compared to classifier-free guidance and classifier guidance, CTRL demonstrates superior performance in controllability for image generation.
summary: This paper proposes a novel method, CTRL, to address adding conditional control to diffusion models. CTRL consists of three main stages: constructing the augmented model with a pre-trained model, training a classifier as a reward model, and solving a reinforcement learning problem to introduce conditional control in the diffusion model. Compared to classifier-free guidance, CTRL is adaptable to more flexible offline datasets, and compared to classifier-guidance methods, CTRL avoids accumulated inaccuracies caused by predicting $y$ from $x_t$ at each denoising step. This makes CTRL applicable in a broader range of scenarios with a more controllable generation process. Finally, experimental results validate the effectiveness of CTRL in terms of controllability and generation quality in tasks involving compressibility and multi-task generation.
weaknesses: 1. The evaluation metrics used in the paper are insufficient; for instance, the author only verifies the controllability of CTRL without providing results on the image quality generated by CTRL.
2. The experimental tasks are limited. The author only uses compressibility and aesthetic pleasingness as conditional controls. It is necessary to compare more complex conditions, such as sketch, normal map, as suggested in [1].

[1] Adding Conditional Control to Text-to-Image Diffusion Models. ICCV, 2023.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
questions: In the error analysis section, the paper briefly discusses three sources of error. Could these three types of errors be more precisely quantified?
soundness: 3
strengths: 1. The using of reinforcement learning to add conditional controls to diffusion models is novel. By reframing conditional generation as an RL problem, this approach leverages optimal policy learning to achieve conditional sampling without needing complex data dependencies. This use of RL significantly improves sample efficiency and offers a new, effective pathway for conditional generation, especially beneficial for complex, multi-condition tasks.
2. The comparison with prior work is comprehensive. The authors thoroughly contrast their RL-based CTRL method with traditional methods such as classifier guidance and classifier-free guidance, highlighting where CTRL overcomes limitations in sample efficiency, dataset requirements, and control precision. By detailing the theoretical and practical distinctions, such as the use of conditional independence to simplify dataset construction, the paper effectively demonstrates how CTRL builds upon and enhances past work, providing readers with a well-rounded understanding of its innovations.
3. Experimental results show that CTRL not only meets target conditions more accurately but also maintains high performance across diverse tasks with fewer data dependencies, underscoring its efficiency and robustness in practical applications.
summary: The paper introduces CTRL to enhance pre-trained diffusion models by integrating additional conditional controls through reinforcement learning. CTRL reformulates the conditional generation as an RL problem, where the reward function is derived from the conditional likelihood of labels given data, and the KL divergence from the pre-trained model acts as a regularizer. This method has key advantages over traditional guidance techniques, as it can use pairs of samples instead of triplets by leveraging conditional independence. CTRL is empirically validated on image generation tasks, demonstrating superior performance in single-task and multi-task settings, including cases requiring compositional controls, like generating images with specific compressibility and aesthetic properties.
weaknesses: 1. While the core RL formulation and the connection to conditional diffusion are well-explained, the theoretical framework could be strengthened with a deeper exploration of the method's convergence properties and guarantees. 
2. The experiments primarily compare CTRL with two standard methods: classifier guidance and classifier-free guidance. Including additional baseline methods, such as more recent variants of guided diffusion techniques or alternative conditional generation approaches, would provide a clearer picture of CTRL’s advantages and limitations.

---

confidence: 2
contribution: 3
initial_rating: 5
presentation: 3
questions: 1. Obviously, CTRL is a splitter-based method, but why is the experimental verification presented as the main comparison with the method without a classifier? Is this unfair? Please explain the considerations behind this.
2. Are there comparisons with more advanced conditional generation models?
3. The authors claim in Appendix G.2 that the proposed method can be adapted to any off-the-shelf RL algorithm, which I take with a grain of salt. There are many kinds of RL algorithms, and the two examples given by the author are both model-free on-policy methods. Did the author take into account offline RL algorithms, model-based RL algorithms, etc.? Is there a more detailed explanation, analysis, or experiment on this?
4. If possible, it is recommended to include a citation to the paper proposing the PPO method [1].

[1] Schulman J, Wolski F, Dhariwal P, et al. Proximal policy optimization algorithms[J]. arXiv preprint arXiv:1707.06347, 2017.
soundness: 3
strengths: 1. This paper builds a good theoretical framework for RL-based conditional control to diffusion model.
2.  As far as I can see, the idea of using reinforcement learning to fine-tune the generation conditions of diffusion models is interesting.
3. The source code is attached in supplementary material, and I believe it is easy to reproduce.
summary: This paper presents a conditional image generation method based on the diffusion model conditions with classifier guidance. Different from other methods of conditional generation of diffusion model in classifier guidance, this paper uses reinforcement learning to optimize conditional image generation of the diffusion model to maximize the log-likelihood probability of conditional generation as the reward function. It eliminates the need to train the classifier according to the intermediate diffusion frame and reduces the difficulty of training the classifier.

Note that as a researcher in reinforcement learning, I pay more attention to the problems related to RL in this paper. And I'm not familiar with image conditional generation.
weaknesses: 1. The experimental part of the paper is very limited, only comparing the two methods in 2022. However, the field of conditional diffusion generation is changing rapidly, and this paper lacks a comparison with the latest methods. This paper does not perform an ablation analysis of the proposed method.
2. The author's specific methods of using RL are not very clear. Only the reward function is introduced, but how to formalize the problem into MDP, the state, the action, the transfer function, and so on are not introduced.

Meta-review: This paper outlines a way to fine-tune a diffusion model, allowing conditioning on additional information. The proposed method is pragmatic because it relies on standard implementations of fine-tuning, based on reinforcement learning.  The authors demonstrate that the proposed method is superior to classifier-free methods (at the additional cost of training a classifier). 

The strengths of the paper are: (1) novel, distinct perspective on training conditional models; (2) good comparison with prior work; (3) availability of source code.

Because of those strengths, I lean towards acceptance.

I do have some residual concerns about presentation, which I hope will be addressed in the camera-ready version. Specifically, the abstract is very difficult to parse. For example, RL is mentioned in a way that suggests that classifiers and KL constraints are an integral part of it (in fact, both can be useful for RL in certain settings, but they are not an integral part of vanilla RL). Also, for some strange reason, the formalisation of the task as an MDP is deferred till the appendix.

Decision: Accept (Poster)