# Adding Conditional Control to Diffusion Models
with Reinforcement Learning

Yulai Zhao
Princeton University
<EMAIL_ADDRESS>
&Masatoshi Uehara∗
Genentech
<EMAIL_ADDRESS>
&Gabriele Scalia
Genentech
<EMAIL_ADDRESS>
&Tommaso Biancalani
Genentech
<EMAIL_ADDRESS>
&Sergey Levine
University of California, Berkeley
<EMAIL_ADDRESS>
&Ehsan Hajiramezanali‡
Genentech
<EMAIL_ADDRESS>
Equal contribution.This work is done during an internship at
Genentech.Corresponding authors.

###### 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 (Conditioning pre-Trained
diffusion models with Reinforcement Learning), 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., 2015; Ho et al., 2020). 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., 2021) and Stable Diffusion (Rombach et al., 2022), $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. (2023)), because it enables
the extension of generative capabilities with new conditional variables
without requiring retraining from scratch. Currently, classifier-free guidance
(Ho et al., 2020) is a prevailing approach for incorporating conditional
controls into diffusion models, and it has proven successful in computer
vision (Zhang et al., 2023; Zhao et al., 2024). 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., 2023), which is often not feasible in many scenarios,
e.g. drug discovery (Huang et al., 2021).

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., 2023; Fan et al.,
2023), 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, 2021) |  No |  Yes |  Yes
Reconstruction guidance (e.g. (Ho et al., 2022), (Chung et al., 2022), (Han et al., 2022)) |  No |  No |  Yes
Classifier-free guidance (Ho and Salimans, 2022) |  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, 2021; Song et al., 2020) 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., 2022; Song et al.,
2022). Table 1 summarizes the main features of the proposed algorithm compared
to existing methods.

##  2 Related Works

##### Classfier guidance.

Dhariwal and Nichol (2021); Song et al. (2020) 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., 2022; Han et al., 2022; Chung et
al., 2022; Finzi et al., 2023; Bansal et al., 2023). 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, 2022) 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., 2021; Saharia et
al., 2022; Rombach et al., 2022). 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.
(2023); Xie et al. (2023). 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., 2023; Wu et al., 2023),
reinforcement learning (Black et al., 2023; Fan et al., 2023; Uehara et al.,
2024), and control-based techniques (Clark et al., 2023; Xu et al., 2023;
Prabhudesai et al., 2023; Uehara et al., 2024). 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 SDE111To 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., 2024).:

$\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 datasets222For 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$ 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, 2021; Bansal et al., 2023; Epstein et al.,
2023), 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, 2021;
Ho and Salimans, 2022; Nichol et al., 2021; Saharia et al., 2022; Rombach et
al., 2022). 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, 2021; Ho
and Salimans, 2022). 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., 2024, Theorem 1).

###  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, 2022) 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.,
2023), ControlNet (Zhang et al., 2023), DiffFit (Xie et al., 2023)). 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., 2021; Yellapragada et al., 2024; Giannone et al., 2024).

####  3.2.2 Classfier Guidance

Classifier guidance (Dhariwal and Nichol, 2021; Song et al., 2020) 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$ 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, 2023).

##### Reconstruction guidance.

To mitigate these issues, several studies (Ho et al., 2022; Han et al., 2022;
Chung et al., 2022; Guo et al., 2024)333We 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., 2022).

##  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$ to $T$. In
the next section, we explain how to solve (5) in practice.

###  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.

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

1: Input: Pre-trained model with a drift coefficient $f^{\mathrm{pre}}$,
Offline data $\mathcal{D}=\\{c^{(i)},x^{(i)},y^{(i)}\\}$, Exploratory
distribution $\Pi\in\Delta(\mathcal{C}\times\mathcal{Y})$

2: Construct an augmented model $g(t,c,y,x;\psi)$.

3: Train a classifier $\hat{p}(y|x,c)$ to approximate $p^{\diamond}(y|x,c)$
from the offline data $\mathcal{D}$

4: Fine-tune the diffusion model by solving the following RL problem (e.g.
using Algorithm 2):
$\displaystyle\hat{\psi}=\mathop{\mathrm{argmax}}_{\psi}\mathbb{E}_{\begin{subarray}{c}(c,y)\sim\Pi(c,y)\\\
x_{0:T}\sim\mathbb{P}^{g}(\cdot|c,y;\psi)\end{subarray}}\left[\gamma\mathbf{\log{}}\hat{p}(y|x_{T},c)-\frac{1}{2}\int_{0}^{T}\frac{\|f^{\mathrm{pre}}(s,c,x_{s})-g(s,c,y,x_{s};\psi)\|^{2}}{\sigma^{2}(s)}ds\right]$
where $\mathbb{P}^{g}(\cdot|c,y;\psi)$ is an distribution induced by the SDE
with a parameter $\psi$.

5: Output: $dx_{t}=g(t,c,y,x_{t};\hat{\psi})dt+\sigma(t)dw_{t}$

##### 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.

Algorithm 2 Direct back-propagation for conditioning

1: Input Batch size $n$, Learning rate $\eta$, Discretization step $\Delta t$,
Exploratory distribution $\Pi\in\Delta(\mathcal{C}\times\mathcal{Y})$.

2: Itinialize: $\psi=[\\{\theta^{\mathrm{pre}}\\}^{\top},\mathbf{0}^{\top}]$

3: for $i\leftarrow 1$ to $S$ do

4: We obtain $n$ trajectories  $\displaystyle\textstyle\\{X^{\langle
k\rangle}_{0},\cdots,X^{\langle k\rangle}_{T}\\}_{k=1}^{n},\\{Z^{\langle
k\rangle}_{0},\cdots,Z^{\langle k\rangle}_{T}\\}_{k=1}^{n}.$ following
$(C^{\langle k\rangle},Y^{\langle k\rangle})\sim\Pi(\cdot),X^{\langle
k\rangle}_{0}\sim\mathcal{N}(0,I_{d}),Z_{0}=0$, and  $\displaystyle X^{\langle
k\rangle}_{t}=X^{\langle k\rangle}_{t-1}+g(t-1,C^{\langle k\rangle},Y^{\langle
k\rangle},X^{\langle k\rangle}_{t-1};\psi_{i})\Delta t+\sigma(t)(\Delta
w_{t}),\,\quad\Delta w_{t}\sim\mathcal{N}(0,(\Delta t)^{2}),$ $\displaystyle
Z^{\langle k\rangle}_{t}=Z^{\langle k\rangle}_{t-1}+\frac{\|g(t-1,C^{\langle
k\rangle},Y^{\langle k\rangle},X^{\langle
k\rangle}_{t-1};\psi_{i})-f^{\mathrm{pre}}(t-1,C^{\langle k\rangle},X^{\langle
k\rangle}_{t-1};\theta^{(i)})\|^{2}}{2\sigma^{2}(t-1)}\Delta t.$

5: Update a parameter:
$\displaystyle\psi_{i+1}=\psi_{i}+\eta\nabla_{\psi}\left\\{\frac{1}{n}\sum_{k=1}^{n}\left[\gamma\log{\hat{p}(Y^{\langle
k\rangle}|X^{\langle k\rangle}_{T},C^{\langle k\rangle})}-Z^{\langle
k\rangle}_{T}\right]\right\\}\bigg{|}_{\psi=\psi_{i}},$

6: end for

7: Output: Parameter $\psi_{S}$

##### 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. (2023); Fan
et al. (2023), 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., 2023; Prabhudesai et al., 2023), 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.,
2021) modules instead of the full diffusion weights, (b) employing gradient
checkpointing (Gruslys et al., 2016; Chen et al., 2016) to conserve memory,
and (c) randomly truncating gradient back-propagation to avoid computing
through all diffusion steps (Clark et al., 2023; Prabhudesai et al., 2023).

######  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. (2023); Fan et al. (2023). 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.

###  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, 2019).

##### 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, 2021). 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$ 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., 2022)],
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.

###  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., 2022) 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., 2023). 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.

##### 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.

##### 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

(a) Histogram of generations’ labels from $p^{\mathrm{pre}}$.

(b) Confusion matrix for DPS.

Figure 1: Statistics of generations from the pre-trained model and the DPS
baseline are shown. In Figure 1(a), we show the histogram for $1024$ images
generated by the pre-trained model (i.e., $p^{\mathrm{pre}}$). We observe that
the samples are distinctly biased, e.g. samples in $Y=3$ are rare, which
renders generating samples given this condition challenging. Next, Figure 1(b)
illustrates the confusion matrix for samples generated by DPS given four
conditions. For each condition, $128$ samples are generated and are evaluated.
We find that DPS struggles to generate samples accurately when conditioned on
intermediate labels.

We compare CTRL with DPS (Chung et al., 2022), 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., 2022) 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

(a) Training curve

| Accuracy $\uparrow$ | Macro F1 score $\uparrow$
---|---|---
DPS | $0.45$ | $0.44$
CTRL (Ours) | $\mathbf{1.0}$ | $\mathbf{1.0}$

(b) Evaluation of conditional generations

(c) Generated images

Figure 2: Results for conditioning on compressibility are presented. Recall
that conditions are established as $Y\in[0,1,2,3]$ to indicate four
compressibility levels with $Y=3$ being the most compressible. Figure 2(a)
illustrates an evaluation curve showing the mean compressibilities of the
generated samples given different $Y$. It is evident that CTRL effectively
fine-tunes the diffusion model, aligning the compressibilities with each
specific condition. Table 2(b) provides accuracies and macro F1 scores for
CTRL and DPS. Additionally, Figure 2(c) showcases images generated by our
approach.

We start by conditioning generations on their file sizes, specifically
focusing on compressibility 444It is important to note that, unlike standard
tasks in classifier guidance (Chung et al., 2022), this score is non-
differentiable w.r.t. images.. In accordance with Black et al. (2023), 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.

(a) Histogram of $1024$ images generated by the pre-trained model.

| Task | Accuracy $\uparrow$ | Macro F1 $\uparrow$
---|---|---|---
DPS | CP | $0.61$ | $0.55$
| AS | $0.66$ | $0.62$
CTRL | CP | $\mathbf{0.94}$ | $\mathbf{0.94}$
| AS | $\mathbf{0.93}$ | $\mathbf{0.92}$

(b) Evaluation of conditional generations

Figure 3:  The results of multi-task generations conditioned on the
composition of compressibility and aesthetic score are presented. Recall that
we determine the conditions as follows: $Y=0$ : low CP and low AS; $Y=1$ :
high CP and low AS; $Y=2$ : low CP and high AS; $Y=3$ : high CP and high AS.
From Figure 3(a), it can be seen that the pre-trained model has very few
generations with labels $Y=1$ and $Y=3$ (i.e., with high CP), thus generations
conditioned on these labels are hard. In Table 3(b), we provide the accuracy
and macro F1 score for each task. We find that CTRL outperforms DPS on both
tasks. Interestingly, our method is capable of synthesizing rare samples from
the pre-trained model due to the benefit of RL-instructed fine-tuning.
Generated images from CTRL are displayed in Figure 3(c).

(c) Generated images

###  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., 2023; Fan et al., 2023;
Uehara et al., 2024), we employ an aesthetic scorer implemented as a linear
MLP on top of the CLIP embeddings (Radford et al., 2021), 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.

## References

  * Anderson (1982) Anderson, B. D. (1982).  Reverse-time diffusion equation models.  Stochastic Processes and their Applications 12(3), 313–326.
  * Bansal et al. (2023) Bansal, A., H.-M. Chu, A. Schwarzschild, S. Sengupta, M. Goldblum, J. Geiping, and T. Goldstein (2023).  Universal guidance for diffusion models.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 843–852.
  * Black et al. (2023) Black, K., M. Janner, Y. Du, I. Kostrikov, and S. Levine (2023).  Training diffusion models with reinforcement learning.
  * Brooks et al. (2023) Brooks, T., A. Holynski, and A. A. Efros (2023).  Instructpix2pix: Learning to follow image editing instructions.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 18392–18402.
  * Chen et al. (2016) Chen, T., B. Xu, C. Zhang, and C. Guestrin (2016).  Training deep nets with sublinear memory cost.  arXiv preprint arXiv:1604.06174.
  * Chung et al. (2022) Chung, H., J. Kim, M. T. Mccann, M. L. Klasky, and J. C. Ye (2022).  Diffusion posterior sampling for general noisy inverse problems.  arXiv preprint arXiv:2209.14687.
  * Clark et al. (2023) Clark, K., P. Vicol, K. Swersky, and D. J. Fleet (2023).  Directly fine-tuning diffusion models on differentiable rewards.  arXiv preprint arXiv:2309.17400.
  * Dhariwal and Nichol (2021) Dhariwal, P. and A. Nichol (2021).  Diffusion models beat gans on image synthesis.  Advances in neural information processing systems 34, 8780–8794.
  * Epstein et al. (2023) Epstein, D., A. Jabri, B. Poole, A. Efros, and A. Holynski (2023).  Diffusion self-guidance for controllable image generation.  Advances in Neural Information Processing Systems 36, 16222–16239.
  * Fan et al. (2023) Fan, Y., O. Watkins, Y. Du, H. Liu, M. Ryu, C. Boutilier, P. Abbeel, M. Ghavamzadeh, K. Lee, and K. Lee (2023).  Dpok: Reinforcement learning for fine-tuning text-to-image diffusion models.  arXiv preprint arXiv:2305.16381.
  * Finzi et al. (2023) Finzi, M. A., A. Boral, A. G. Wilson, F. Sha, and L. Zepeda-Núñez (2023).  User-defined event sampling and uncertainty quantification in diffusion models for physical dynamical systems.  In International Conference on Machine Learning, pp. 10136–10152. PMLR.
  * Giannone et al. (2024) Giannone, G., A. Srivastava, O. Winther, and F. Ahmed (2024).  Aligning optimization trajectories with diffusion models for constrained design generation.  Advances in Neural Information Processing Systems 36.
  * Gruslys et al. (2016) Gruslys, A., R. Munos, I. Danihelka, M. Lanctot, and A. Graves (2016).  Memory-efficient backpropagation through time.  Advances in neural information processing systems 29.
  * Guo et al. (2017) Guo, C., G. Pleiss, Y. Sun, and K. Q. Weinberger (2017).  On calibration of modern neural networks.  In International conference on machine learning, pp. 1321–1330. PMLR.
  * Guo et al. (2024) Guo, Y., H. Yuan, Y. Yang, M. Chen, and M. Wang (2024).  Gradient guidance for diffusion models: An optimization perspective.  arXiv preprint arXiv:2404.14743.
  * Han et al. (2022) Han, X., H. Zheng, and M. Zhou (2022).  Card: Classification and regression diffusion models.  Advances in Neural Information Processing Systems 35, 18100–18115.
  * Ho et al. (2020) Ho, J., A. Jain, and P. Abbeel (2020).  Denoising diffusion probabilistic models.  Advances in neural information processing systems 33, 6840–6851.
  * Ho and Salimans (2022) Ho, J. and T. Salimans (2022).  Classifier-free diffusion guidance.  arXiv preprint arXiv:2207.12598.
  * Ho et al. (2022) Ho, J., T. Salimans, A. Gritsenko, W. Chan, M. Norouzi, and D. J. Fleet (2022).  Video diffusion models.  Advances in Neural Information Processing Systems 35, 8633–8646.
  * Hu et al. (2021) Hu, E. J., Y. Shen, P. Wallis, Z. Allen-Zhu, Y. Li, S. Wang, L. Wang, and W. Chen (2021).  Lora: Low-rank adaptation of large language models.  arXiv preprint arXiv:2106.09685.
  * Huang et al. (2021) Huang, K., T. Fu, W. Gao, Y. Zhao, Y. H. Roohani, J. Leskovec, C. W. Coley, C. Xiao, J. Sun, and M. Zitnik (2021).  Therapeutics data commons: Machine learning datasets and tasks for drug discovery and development.  In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 1).
  * Lee et al. (2023) Lee, K., H. Liu, M. Ryu, O. Watkins, Y. Du, C. Boutilier, P. Abbeel, M. Ghavamzadeh, and S. S. Gu (2023).  Aligning text-to-image models using human feedback.  arXiv preprint arXiv:2302.12192.
  * Li and van der Schaar (2023) Li, Y. and M. van der Schaar (2023).  On error propagation of diffusion models.  In The Twelfth International Conference on Learning Representations.
  * Loshchilov and Hutter (2019) Loshchilov, I. and F. Hutter (2019).  Decoupled weight decay regularization.  In International Conference on Learning Representations.
  * Murray et al. (2012) Murray, N., L. Marchesotti, and F. Perronnin (2012).  Ava: A large-scale database for aesthetic visual analysis.  In 2012 IEEE conference on computer vision and pattern recognition, pp. 2408–2415. IEEE.
  * Nichol et al. (2021) Nichol, A., P. Dhariwal, A. Ramesh, P. Shyam, P. Mishkin, B. McGrew, I. Sutskever, and M. Chen (2021).  Glide: Towards photorealistic image generation and editing with text-guided diffusion models.  arXiv preprint arXiv:2112.10741.
  * Prabhudesai et al. (2023) Prabhudesai, M., A. Goyal, D. Pathak, and K. Fragkiadaki (2023).  Aligning text-to-image diffusion models with reward backpropagation.  arXiv preprint arXiv:2310.03739.
  * Radford et al. (2021) Radford, A., J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, P. Mishkin, J. Clark, G. Krueger, and I. Sutskever (2021).  Learning transferable visual models from natural language supervision.  arXiv preprint arXiv:2103.00020.
  * Ramesh et al. (2021) Ramesh, A., M. Pavlov, G. Goh, S. Gray, C. Voss, A. Radford, M. Chen, and I. Sutskever (2021).  Zero-shot text-to-image generation.  In International conference on machine learning, pp. 8821–8831. Pmlr.
  * Rogers and Williams (2000) Rogers, L. C. G. and D. Williams (2000).  Diffusions, Markov processes and martingales: Volume 2, Itô calculus, Volume 2.  Cambridge university press.
  * Rombach et al. (2022) Rombach, R., A. Blattmann, D. Lorenz, P. Esser, and B. Ommer (2022).  High-resolution image synthesis with latent diffusion models.  In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 10684–10695.
  * Saharia et al. (2022) Saharia, C., W. Chan, S. Saxena, L. Li, J. Whang, E. L. Denton, K. Ghasemipour, R. Gontijo Lopes, B. Karagol Ayan, T. Salimans, et al. (2022).  Photorealistic text-to-image diffusion models with deep language understanding.  Advances in neural information processing systems 35, 36479–36494.
  * Shreve et al. (2004) Shreve, S. E. et al. (2004).  Stochastic calculus for finance II: Continuous-time models, Volume 11.  Springer.
  * Sohl-Dickstein et al. (2015) Sohl-Dickstein, J., E. Weiss, N. Maheswaranathan, and S. Ganguli (2015).  Deep unsupervised learning using nonequilibrium thermodynamics.  In International conference on machine learning, pp. 2256–2265. PMLR.
  * Song et al. (2020) Song, J., C. Meng, and S. Ermon (2020).  Denoising diffusion implicit models.  arXiv preprint arXiv:2010.02502.
  * Song et al. (2022) Song, J., A. Vahdat, M. Mardani, and J. Kautz (2022).  Pseudoinverse-guided diffusion models for inverse problems.  In International Conference on Learning Representations.
  * Song et al. (2020) Song, Y., J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole (2020).  Score-based generative modeling through stochastic differential equations.  arXiv preprint arXiv:2011.13456.
  * Uehara et al. (2024) Uehara, M., Y. Zhao, K. Black, E. Hajiramezanali, G. Scalia, N. L. Diamant, A. M. Tseng, T. Biancalani, and S. Levine (2024).  Fine-tuning of continuous-time diffusion models as entropy-regularized control.  arXiv preprint arXiv:2402.15194.
  * Uehara et al. (2024) Uehara, M., Y. Zhao, K. Black, E. Hajiramezanali, G. Scalia, N. L. Diamant, A. M. Tseng, S. Levine, and T. Biancalani (2024).  Feedback efficient online fine-tuning of diffusion models.  arXiv preprint arXiv:2402.16359.
  * Uehara et al. (2024) Uehara, M., Y. Zhao, E. Hajiramezanali, G. Scalia, G. Eraslan, A. Lal, S. Levine, and T. Biancalani (2024).  Bridging model-based optimization and generative modeling via conservative fine-tuning of diffusion models.  arXiv preprint arXiv:2405.19673.
  * Wainwright (2019) Wainwright, M. J. (2019).  High-dimensional statistics: A non-asymptotic viewpoint, Volume 48.  Cambridge university press.
  * Wu et al. (2023) Wu, X., K. Sun, F. Zhu, R. Zhao, and H. Li (2023).  Better aligning text-to-image models with human preference.  arXiv preprint arXiv:2303.14420.
  * Xie et al. (2023) Xie, E., L. Yao, H. Shi, Z. Liu, D. Zhou, Z. Liu, J. Li, and Z. Li (2023).  Difffit: Unlocking transferability of large diffusion models via simple parameter-efficient fine-tuning.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 4230–4239.
  * Xu et al. (2023) Xu, J., X. Liu, Y. Wu, Y. Tong, Q. Li, M. Ding, J. Tang, and Y. Dong (2023).  Imagereward: Learning and evaluating human preferences for text-to-image generation.  arXiv preprint arXiv:2304.05977.
  * Yellapragada et al. (2024) Yellapragada, S., A. Graikos, P. Prasanna, T. Kurc, J. Saltz, and D. Samaras (2024).  Pathldm: Text conditioned latent diffusion model for histopathology.  In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision, pp. 5182–5191.
  * Zhang et al. (2023) Zhang, L., A. Rao, and M. Agrawala (2023).  Adding conditional control to text-to-image diffusion models.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 3836–3847.
  * Zhang et al. (2023) Zhang, Z., L. Han, A. Ghosh, D. N. Metaxas, and J. Ren (2023).  Sine: Single image editing with text-to-image diffusion models.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 6027–6037.
  * Zhao et al. (2024) Zhao, S., D. Chen, Y.-C. Chen, J. Bao, S. Hao, L. Yuan, and K.-Y. K. Wong (2024).  Uni-controlnet: All-in-one control to text-to-image diffusion models.  Advances in Neural Information Processing Systems 36.

##  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 555In 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$ 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., 2020), 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, 1982), 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)

∎

##### 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., 2004), 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].$

∎

###  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$.

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$.

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, 2019) 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., 2023; Prabhudesai et al., 2023).

##### 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., 2023), in
practice it is often desired to make minimal changes to the architecture of
large diffusion models, e.g., Stable Diffusion (Rombach et al., 2022) 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., 2020). 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. (2023); Prabhudesai et al. (2023): (1)
Fine-tuning low-rank adapter (LoRA) modules (Hu et al., 2021) instead of
tuning the original diffusion weights, and (2) Gradient checkpointing for
computing partial derivatives on demand (Gruslys et al., 2016; Chen et al.,
2016). 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., 2012), 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., 2021). To train the classifier, we use the full
AVA dataset (Murray et al., 2012) 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. (2017).

###  D.2 Additional Results

##### More generated images.

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

Figure 4: More images generated by CTRL in the compressibility task.  Figure 5: More images generated by CTRL in the multi-task conditional generation.  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.43$ |  $-110$
$0$ |  $5$ |  $0.64$ |  $-157.4$
$0$ |  $7.5$ |  $0.62$ |  $-148.2$
$0$ |  $10$ |  $0.52$ |  $-156.5$
$0$ |  $20$ |  $0.66$ |  $-152.5$
$0$ |  $50$ |  $0.69$ |  $-151.6$
$1$ |  $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$ |  $-194.0$
$2$ |  $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$ |  $-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. 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. 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. 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)