Adding Conditional Control to Diffusion Models with Reinforcement Learning
Authors: Yulai Zhao
Masatoshi Uehara
Gabriele Scalia
Sunyuan Kung
Tommaso Biancalani
Sergey Levine
Ehsan Hajiramezanali
Venue: ICLR 2025 Poster (2025)
Keywords: diffusion models
conditional generation
fine-tuning
reinforcement learning
Primary area: generative models
TL;DR: This work introduces a provable RL-based fine-tuning approach for conditioning pre-trained diffusion models on additional controls.

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 such controls using an offline dataset comprising inputs and 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. Our method, **CTRL** (**C**onditioning pre-**T**rained diffusion models with **R**einforcement **L**earning), produces soft-optimal policies that maximize the abovementioned reward functions. We formally demonstrate that our method enables sampling from the conditional distribution with additional controls during inference.
Our RL-based approach offers several advantages over existing methods. Compared to classifier-free guidance,
it improves sample efficiency and can greatly simplify dataset construction by leveraging conditional independence between the inputs and additional controls. Additionally, unlike classifier guidance, it eliminates the need to train classifiers from intermediate states to additional controls.
The code is available at https://github.com/zhaoyl18/CTRL.

Published as a conference paper at ICLR 2025
ADDING CONDITIONAL CONTROL TO DIFFUSION MOD- ELS WITH REINFORCEMENT LEARNING
Yulai Zhao
Masatoshi Uehara"
Princeton University
Genentech uehara masatmsshi@eeeeeee com
yulaiz@pr nceton.
Gabriele Scalia
Sunyuan Kung Princeton University
Tommaso Biancalani Genentech biancalt@gene. com
Genentech scaliagegene.com
kung@princetonn.c edu
Sergey Levine University of Californa, Berkeley sergey wew.ne@eerkeeey......ooe edu
Ehsan Hajiramezanali
Genentech
ha jaairamm@genr gen.e. com
ABSTRACT
Diffusion models are powerful generative models that allow for precise control over the characteristics of the generated samples. While these diffusion models trained large datasets have achieved success, there is often need to introduce additional controls downstream fine-tuning processes, treating these powerful models as pre-trained diffusion models. This work presents novel method based rein- forcement learning (RL) to add such controls using an offline dataset comprising inputs and labels. We formulate this task as an RL problem, with the classifier learned from the offline dataset and the KI. divergence against pre-trained models serving the reward functions. Our method, CTRL (Conditioning pre-Trained dif- fusion models with Reinforcement Learning), produces soft-optimal policies that maximize the abovementioned reward functions. We formally demonstrate that our method enables sampling from the conditional distribution with additional controls during inference. Our RL-based approach offers several advantages over existing methods. Compared to classifier- free guidance, it improves sample efficiency and can greatly simplify dataset construction by leveraging conditional independence between the inputs and additional controls. Additionally, unlike classifier guidance, it eliminates the need to train classifiers from intermediate states to additional con- trols. The code is available at attttps: //gtttb.tttttt////// com/ zhacyl18 CTRL
INTRODUCTION
Diffusion models have emerged as effective generative models for capturing intricate distribu- tions (Sohl-Dickstein et al. 2015 Ho et al. 2020). Their capabilities are further enhanced by building conditional diffusion models p p(z c). For instance, in text-to-image generative models like DALL- E (Ramesh et al. 202I) and Stable Diffusion (Rombach al. 2022), C is a prompt, and X is the image generated according to this prompt. While diflusion models trained on extensive datasets have shown remarkable success, additional controls often need be incorporated during downstream fine-tuning when treating these powerful models as pre-trained diffusion models.
In this work, goal is to incorporate Specifically, given access to a large pre-trained model capable of modeling p(z[c) trained on extensive datasets, we aim condition it additional random variable y V, thereby creating generative model p(z) c, y). To accomplish this, we utilize the pre-trained model and an offline dataset consisting
conditional controls into pre- pre-trained diffusion models.
"Equal contribution.
This work is done during an internship at Genentech.
"oorrsponniig authors.

---

Published as a conference paper at ICLR 2025
Table 1: Comparison with existing approaches. Unlike classifier guidance and its variations, our method involves directly fine-tuning pre-trained models. We avoid the need to learn mapping from It >> y or rely on heuristic approximations. Additionally, while classifier guidance always demands triplets {c, x. v). our approach can leverage conditional independence (y 1 c/r) and only necessitate pairs {ar y). simplifying the construction of the offline dataset.
Leveraging conditional independence
Need to learn 3t + y
Methods
Fine-tuning
Classifier guidance (Dhariwal and Nichol 2021
No
Yes
Yes
Reconstruction guidance (e.g- HH al.
2022);
No
No
Yes
Chung et al. (2022); Han et al. (2022))
Classther-Iree guidance (Ho and Salımans
2022)
Yes
No
No
CTRL((urm)
Yes
No
Yes
of triplets {c, 2. y). This scenario is important, as highlighted in the existing literature on computer vision (c.g., Zhang et al. (2023)) because it enables the extension of generative capabilities with new conditional variables without requiring retraining from scratch. Currently, classifier-fres guidance (Ho et al. (020) is 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 extend well to other challenging problems, especially when large offline datasets are unavailable. Indeed, the success of training conditional diffusion models via classifier- free guidance heavily relies such datasets (Brooks al. 2023), which are often impractical to
obtain. In these scenarios, this method tends to struggle.
In our work, present new approach for adding new conditional controls via reinforcement learning (RL) to further improve sample efficiency. Inspired by recent progress in RL-baased fine- tuning (Black et al. 2023) Fan et al. 2023 Uehara et al. (2024), we frame the conditional generation as an RL problem within a Markov Decision Process (MDP). In this formulation, the reward, which want to maximize, is the (conditional) log-likelihood function log pep |z, c), and the policy, conditioned on (c, y), corresponds to the denoising process at each time step in diffusion model. We formally demonstrate that, by executing the soft-optimal policy, which maximizes the reward log p(y|æ,, c) with KI. Xe.pettyty against the pre-trained model, we can sample from the target conditional distribution p(r|c, y) during inference. Hence, our proposed algorithm, CTRL (Conditioning pre- Trained diffusion models with Reinforcement Learning) consists of three main steps: (1) learning classifier log p(y) x c) (which will serve as our reward function in the MDP) from the offline dataset, augmented diffusion model by adding (trainable) parameters to the pre-trained
(2) constructing model in order to accommodate an additional label y. and (3) learning soft-optimal policy within the aforementioned MDP during fine-tume munitunnng. Our approach is novel as it significantly diverges from classifier-free guidance and distinguishes itself from existing RL-based fine-tuning methods by integrating augmented model in the fine-tuning process support additional controls.
Our novel RL-based approach offers several advantages over existing methods for adding additional controls. Firstly, in contrast classifier-free guidance, which uses offline data to directly model p(zly,c) c). our method leverages offline data to model the simpler distribution p(ylæ, c), improving sample efficiency (in typical scenarios where y is liwer-dimensional than æ). Secondly, in typical scenarios where the additional label y depends solely on z (e.g., the compressibility of an image depends only on the image, not the prompt), our fine-tuning method only requires pairs {a, whereas classifier-free guidance still necessitates triplets {c, x. y) from the offline dataset. This is because the reward function simplifies to log p(g((())) due to the conditional independence y clar, which gives log ply| c) log p()) Furthermore, when the goal is to simultaneously add conditioning controls on two labels, 311 and y2, and both labels only depend on x, our method requires only pairs {xr,9) y1 ) and {x, y2 }. In contrast, classifier-free guidance requires quadruples {c, x, V1, y2). Therefore, in this manner, CTRL can also leverage the compositional nature of the mapping between inputs and additional labels.
Our contributions can be summarized as follows. We propose an Ra-based fine-tuning approach for conditioning pre-trained diffusion models on additional labels. In comparison classifier-free
2

---

Published as a conference paper at ICLR 2025
guidance, 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. Additionally, we establish close connection to classifier guidance (Dhariwal and Nichol 2021; Song et al., 2020), showing that it provides an alternative method for obtaining the afore- mentioned soft-optimal policies (in ideal cases where there are no cesesfiaa/mmodle-iisseeefffieifi errors in the algorithms). Despite this connection, our algorithm addresses common challenges in classifier guidance, such as the need to learn classifiers at multiple noise scales in standard classi- fier guidance and the use of fundamental approximations in some variants to avoid learning these noisy classifiers (Chung al., [022] Song et al.) 2022). Experimentally, we validate the superiority of CTRL over baselines in both single-task and multi-task conditional image generation, such as generating highly aesthetic yet compressible images, where existing methods often struggle. Table T 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 classifier and incorporating its gradients to guide inference (while freezing pre-trained models). However, notable drawback of this technique lies in the classifier's accuracy in predicting from intermediate Ty, resulting in cumulative errors during the diffusion process. To mitigate this issue, several studies propose methods to circumvent it through recon- struction, referred to as reconstruction guidance in this work. Specifically, they employ certain approximations that map intermediate states 34 back the original input space 20. allowing the classifier to be learned solely from Zy to (Ho et al., 2022 Han et al. 2022 |022|Chung et al.) 2022 Finzi
et al. 2023 Bansal et al. 2023). In contrast to these works, our approach focuses fine- men-ttning 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-fre free guidance (Ho and Salimans 2022 is 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 exxttto-image models (Nichol et al. 2021: Saharia et al. 2022] Rombach et al. 2022). While the original research does not explore elassifier-rre free guidance within the scope of fine-tuning pre-trained diffusion models, several subsequent studies address fine-tuning scenarios Zhang et al. (2023); Xic 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 offline dataset construction.
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 al. 2023). reinforcement learning (Black et al.) [023] Fan et al. 2023)
Uehara et al. 2024). and
techniques (Clark et al. 20233] Xu et al. [023 Prabhudesai al. 2023 Uehara 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.
Remark 1. A concurrent study (Denker el al. (025) presents similar key theorem demonstrating that conditioning can be framed as reinforcement learning (RL) problems. However, while their exper- iments rely on denoising score-mate matehing objectives, our work empirically validates the effectiveness of RL objectives.
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 y and C to represent condition spaces and X' to denote the (Fuclidean) sample space. Given the tre-trd
3

---

Published as a conference paper at ICLR 2025
model, which enables us to sample from pa (x)e) : : C A(Y), our goal is to add new conditional controls E y such that we can sample from y))
Pre-trained model and offline dataset. A (continuous-time) pre-trained conditional diffusion
model characterized by the following SDF where frre: :0, T x C x X' IR" a model with parameter 0:
dx = (L, c, 343 dd + (()dwr,
Jo = Cini:
In training diffusion models, the parameter frme is derived by optimizing specific loss function on large datasets P
We refer interested readers to Appendix A for more details on constructing these loss functions. Using the pre-trained model and following the above SDE from 0 to T, sample from pPre (-Ic) for any condition c CE6 C.
To add additional control to pre-trained model, as in many recent works (Dhariwal and Nichol 2021
Bansal et al. 2023; Epstein et al. 2023), we assume access to offline data: D = {c(0) (1) ( ii=
C X X x D We denote the conditional distribution
given r and c by p" (s)lx,c).
Target distribution. Using the pre-trained model and the offline dataset, our goal is to obtain diffusion model such that we can sample from a distribution over C x y - ( as below:
(p), (-|c)
Py (+c,, y) He p (x)c))())))))
where 7 E IR+ denotes the strength of additional guidance and μ 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. 202T; Saharia et al.) 2022; Rombach et al. (2022). Specifically, when 7 = I, this distribution corresponds to the standard conditional distribution p c. y). which a fundamental objective of many conditional generative models (Dhariwal and Nichol, 2021 Ho and Salimans 2022 Moreover, for general y, Py can be
formulated via the following optimization problem:
Py (-c,,)) = argmin (x-g(- (ean log (y c)] KL(g(-le, 14
log
(CX)~A(X)
This relation clear as the objective function equals (L(g(-le, (+,c, up to a constant.
Goal. As discussed, the primary goal of this research is to train a generative model capable of simulating Py (lc,,)) To achieve this, we introduce the following SDE:
dr; = g(t,
xx) dt + (t) dw,, xo
Tini
(3)
where g: 10, 7J x C X + d is an augmented model to add additional controls into pre-trained models. The primary challenge involves leveraging both offline data and pre- pre-trained model weights to train the term g. ensuring that the marginal distribution of zr induced by the SDE (3) accurately approximates P7-
Notation. Let the space of trajectories Fo.r be K.. Conditional on and y, we denote the measure induced by the SDE (3) over IC by (- y). Similarly, we use Pi (-Ic, y) and (-lc, y) to represent the marginal distribution of I, and density dP? (rlc, sd
3.2 EXISTING METHODS
In this subsection, we describe two existing methods that are applicable in our context to achieve the aforementioned goal.
3.2. 1
CLASSFIER-FRES GUIDANCE
Recent works have studied fine-tuning pre-trained models with classifier-free guidance (Brooks et al. 2023 Zhang et al. 2023; Xie et al. 2023). These methods introduce an augmented model g as described in 33, where the weights are initialized from the pre-trained model. Fine-tuning is via minimizing classifier-fire free guidance loss on the offline dataset. While successful in many applications, these methods may struggle in scenarios where offline datasets for new conditions are limited (Huang
et al
2021 Yellapragada et al. 2024 Giannone et al. 2024).
For simplicity, we consider case where the initial distribution is Dirac delta, as in bridge matching.
The extension of our proposal for stochastic distributions remains straightforward (Uchara ct al. 2024).
For notational simplicity, throughout this work, we would often drop grte.
4

---

Published as a conference paper at ICLR 2025
3.2.2 CLASSFIER GUIDANCE
Classifier guidance (Dhariwal and Nichol, 2021 Song et al. 2020) is based on the following result.
Lemma (Doob's 's h-transforms (Rogers and Williams 2000).. For any e E C and y E Y, by evolving according the following SDE from to T: dx = {f (t, c, x4) + o (t) v., log En our (1 (lar.))
= - s d] + (t))du,,
(4)
Additional Vx log logp(yle..e)
the marginal distribution or,
æT, ie., p(rrle, y), is equal the target distribution Py=1 (-1c, (c,y) Here, pure denotes the distribution induced by the pre-trained diffusion model ①
This lemma suggests that simulate the target distribution (2), we only need construct SDE (4). However, practical issues arise: first, training classifier p(y) D2, e) requires extensive data at each timestep, which is cumbersome with large pre-trained models. Furthermore, accumulated inaccuracies in drift estimates may lead to poor performance (Li and van der Schaar 2023).
Reconstruction guidance. To mitigate these issues, several studies propose approximate p(ylz 2, c) directly via reconstruction (Ho et al. 2022: Han et al. 2022 Chung et al. 2022 Guo al. 2024 specifically by p(y] e) P ( TT, c C(r(())(pae (t, 6 der R plyler C), c). where (T,, c) the (expected) denoised sample given Tt. i.e., Fr (3x,, c) E[ cj. Given such
an approximation, we only need to learn p (yl æT, e) from data. However, this approximation may become imprecise when P(ærlæz lz, c) is noisy is difficult to estimate reliably (Chung et al. 2022
4
CONDITIONING
DIFFUSION MODELS WITH RL
This section provides details on how our method solves the aforementioned goal with methodolog- ical motivations. We begin with key observation: the conditioning problem be effectively conceptualized as an RL problem 4] Building upon this insight, we illustrate our main algorithm.
4.1 CONDITIONING AS RI
Recall that our objective is to learn a drift term in (3) so that the induced marginal distribution at T (i.e., pf)) closely matches our target distribution py. To achieve this, we first formulate the problem via the following minimization:
argmin (e.
With some algebra, we can show that the above optimization problem equivalent to the following:
/'' [ (s, c, x.) g(8 c, y, x1 ds
argmin (x -7 log (ulo, 20 c)
(²
Here, recall that ps the measure induced by SDE
(3) with a drift coefficient g. Based on this
observation, we derive the following theorem.
Theorem (Conditioning as RL). Consider the following RL RL. problem:
1 1/ f (s, C. - g(s. c, U. v,
g" = argmax ((a)) ParropP (-laa))
n >logp' (ylep,yyrrerrr c)
2
02(s)
ds
g
(5)
where IT A(C x V). Significantly, the marginal distribution pi matches our target distribution: V(c, y) SSupp(I)); pf. (-Ic, p-lc,))
Proofs are deferred to Appendix D This theorem demonstrates that, after obtaining the optimal drift g" by solving the RL problem (5), we can sample from the target distribution p- y) by following SDE (3) from time 0 to T. In the next section, we explain how to solve 5 in practice.
4.2 ALGORITHM
Theoretically inspired by Theorem we introduce Algorithm I It consists of three steps.
We categorize them as reconstruction guidance methods for simplicity. note there are many variants. 'For more details of the RI. formulation, such state space, action space, and transition function, please
refer (Cchara al. 2024
5

---

Published as a conference paper at ICLR 2025
Algorithm 1 Conditioning pre- pre-Trained diffusion models with Reinforcement Learning (CTRL)
1: Input: Pre-trained model with a drift coefficient frre, Offline data D = { xx (0), y Ex-
ploratory distribution II AC 3)
2: Construct an augmented model g(t, c, y.
3: Train classifier p[y| c) approximate pº (yla,, from the offiin data DD
4: Fine-tune the diffusion model by solving the following RL problem (e.g, using Algorithm 2):
2 /"'
|| (p.ross,,,, - g(s, c, y, Rsi 1 ds
e = argmax (c.0)~-(e..)) =arP-* (-legr)))
ylog ( c)
o²(s)
where PV (.1c, yi tb) is an distribution induced by the SDE with parameter vA.
5: Output: dx,- = g(t, VI dd + at))dw
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 (4, 0). We introduce an augmented model g(t, y, x; 2 with parameters & = [0 o IT initialized at ننع [ _O]. Here, 1/ structured as a combination of the existing parameters and new parameters é.
Determining the specific architecture of the augmented model involves a tradeoff: adding more new parameters enhances expressiveness but raises computational costs. scenarios where y is discrete with cardinality D*I the most straightforward solution is to instantiate & with a simple linear embedding layer that maps each y E y to corresponding embedding. These embeddings then added to every intermediate output in the diffusion SDE (i.e., 24 in (3). This method
preserves the original structure to the fullest extent while ensuring that all pre-trained weights fully utilized. Experimentally, we observe that this lightweight modification leads to accurate conditional generations for complex conditioning tasks, as shown in /eection 6
Step 2: Training a calibrated classifier with offline data (Line 3). Using a function class F C x ((V),, we perform maximum likelihood estimation (MLE):
p(-12, c) == argmax log r(y( c()). rEF
(6)
For instance, when Y is discrete, this loss reduces to the standard cross-entropy loss. When y is continuous, assuming Gaussian noise, it reduces to regression loss.
Step 3: Planning (Line 4). Equipped with a classifier, we proceed to solve the RL problem (5). which constitutes the core 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, many types of off-the-shelf RL algorithms can be employed for planning. In this work, inspired by (Clark et al. 2023 Prabhudesai et al. 2023), we employ direct back- propagation, which requires differentiable. If classifier is not differentiable,
recommend using PPO-based methods (Schulman et al.) 2017 Black et al.) 2023). Details are deferred to Appendix
Below, we make several remarks regarding implementing CTRL in practice.
Remark (Using classifier free guidance to adjust guidance strength). Throughout the fine-tuning process demonstrated in Algorithm 7] the guidance strength for the additional conditional control
(iee, y) is fixed at a specific (see the target conditional distribution (2). However, we note that during inference, this guidance strength can be adjusted-either increased or decreased-using the classifier-free guidance technique. Details are deferred to Appendix C
Remark 3 (Choice of exploratory distribution II). According to Theorem 1 it is desired to improve the coverage over C y during fine-tuning For example, in practice, y only takes several discrete values, we can sample y E y uniformly from these values as done Section
4.3 SOURCE OF ERRORS IN CTRL
We discuss the potential sources of error that CTRL may encounter, which will be useful for comparison with existing methods in the next section. Additional limitations, such as computational cost, memory complexity, and the choice of guidance strength 7, are discussed in Appendix F

---

Published as a conference paper at ICLR 2025
Statistical error. Statistical errors arise during the training of classifier p(y) c) from offiin data while learning plyy/z, c). typical statistical error is given by:
B(z,c)--orr [l c) - p" (-la, li O(Cap(F)/n))
(7)
where jo E A x C) represents the distribution of offline data, and Cap(F) denotes the size of the function class F Wainwright 2019).
Model-misspeccfication error. Model-misspeciiiaalie errors may occur during the learning of the classifier pº [y] c) and in the augmented model if it fails capture the optimal drift g".
Optimization error. planning step.
Optimization errors may occur during both the classifier training step and the
ADDITIONAL COMPARISONE WITH EXISTING CONDITIONING METHODS
In this section, further clarify the connections and comparisons between algorithm and the existing methods.
5.1 COMPARISON TO CLASSIFIER GUIDANCE
We explore the advantages of CTRL over classifier guidance (Dhariwal and Nichol 2021), while also offering theoretical insights that link the two approaches. Despite their distinct goals classifier guidance is an inference-ting timme technique, whereas our method fine-tunes an augmented diffusion model, there is deep theoretical connection. This link is highlighted by our derivation of the analytical expression for the optimal drift in the RL problem (5) as below.
Lemma 2 (Bridging RL-based conditioning with classifier guidance). The optimal drift term g" for RL problem (5) has the following explicit solution:
c,,,) 9" (t, v, f (t, x1)= + +z) 2: log (p ()) c). Vt E [0,2]
The proof of Lemma 2 is deferred to Appendix D.3 This lemma indicates that when y 1, the optimal drift g" corresponds to the drift term obtained from Doob's h-transform (i.e., Lemma I). which is precise used formula classifier guidance. Despite the link to classifier guidance through Lemma 2] our algorithm is fundamentally different. Classifier guidance requires learning predictor from 3e to y for every E T, leading to accumulated inaccuracies. In contrast, our algorithm directly solves the RL problem 5), avoiding the need for such predictors.
In Section 3.2.2] we explore reconstruction guidance methods that also aim to circumvent predicting y from Ty. These methods propose first mapping I; a denoised estimate (π(Z;) and using this for further computations. However, this approximation can be imprecise, especially over longer time horizons. As shown in (Chung et al. 2022 Theorem 1), inherent approximation errors persist without statistical, modell-mssppecifccttion or optimization errors. In contrast, our algorithm avoids such approximation errors.
5.2 COMPARISON CLASSIFIER-FRES GUIDANCE
We first show how CTRL leverages conditional independence to ease implementation, feature absent classifier-fre free guidance. Finally, we discuss the improvements regarding sample (statistical) efficiency.
5.2.1 LEVERAGING CONDITIONAL INDEPENDENCE, COMPOSITIONALLY VIA CTRL
We discuss two scenarios where our method outperforms the classifier-free approach by exploiting the conditional independence between inputs and additional controls.
Example (Scenario Y 1. C)X), If new condition Y is conditionally independent of an existing condition C' given X, meaning pº (y)2, c) p(y(z)) This allows CTRL to operate efficiently with just (z, y) pairs, avoiding the need (c, x, y) triplets in the offline dataset.
This scenario is common in practice. For example, when using the Stable Diffusion pre-trained model (Rombach et al. 2022). where X an image and C is a text prompt, we may also want Y, such score functions like compressibility, aesthetic score,
to condition the generations color (Black et al. 2023). These scores depend solely on the image and are independent of the prompt, meaning Y 1 C|X. We further explore this scenario our experimental analysis in Section 6.1
7

---

Published as a conference paper at ICLR 2025
Multi-task conditional generation. Multi-task conditional generation is significant challenge. requiring the integration of multiple controls into pre-trained models. In the following example, we show how method can be extended to handle this.
Example 2 (Scenario Yi T Y2 X, C). If two conditions, Y1 and Y2, exhibit conditional independence given X and C, such that poyp(1 y2 (z, c) = plyp (r, + log plyz|z, c). the two classifiers can be trained separately using (c, x. 1/1 and (c, x. U2) triplets. Furthermore, if Y1 and Y2 are also independent of C given X (as in Example 1). the classifiers can be trained solely with (xr, y1) and (x, 1/2) pairs, significantly simplifying datasel construction.
This scenario is also common in practice. For instance, with the Stable Diffusion pre-trained model, where X is an image and is a text prompt, additional attributes like Y1 (compressibility) and Y2 (color) depend only on the image, not the prompt. Thus, we can leverage the conditional independence of 11 and Y2 from C given X to simplify the implementation of CTRL. The effectiveness of CTRL in this context is further validated experimentally in Seccion 6.2)
Can classifier-fres free guidance leverage conditional independence? independence in classifier-fre guidance, which directly models p-, y). uncertain. For instance, when Y T. C as in Example I] our method only requires (z. y) pairs, while classifier-fres free guidance typically needs (r, y) triplets. when Y1 Y2|C, X as in [Example our approach utilizes triplets (c, v, y1) and (c, x. 1/2). However, as far as we are concerned, quadruples (c, y1, 32) are necessary for classifier- free guidance, and acquiring such data at scale could pose a bottleneck.
The applicability of conditional
5.2.2 STATISTICAL EFFICIENCY
We present the rationale for our approach being more sample-efficient than classifier-free guidance. Most importantly, we leverage pre-trained model to sample from (x]cc) which is already trained on large datasets. This allows us to focus only on learning the classifier (y)l,, c) from offline data. As a result, any statistical errors from the offline data affect only the classifier learning step (6). In contrast, classifier-free guidance attempts to model the entire distribution pr(-le, y) directly from offline data. Therefore, our method is more sample-efficient by learning only the necessary components from the offline data.
EXPERIMENTS
We compare CTRL with five baselines: (1) Reconstruction Guidance. It attempts alleviate the approximation error of classifier guidance via reconstruction. (2) Classifier-Frre Free guidance (Ho and
Salimans 2022 (3) SMC (Sequential Monte Carlo). Recent works (Wu al., 2024 Phillips
et al. (024) leverage resampling techniques to approximate distributions in diffusion models across batch of samples (i.e., particle filtering). This method is training-free free. (4) SVDD (Li al. 2024).
method that iteratively selects preferable samples during each diffusion step based on reward signals. For our conditioning task, we use a trained classifier as the reward function. This method is also training- frrennngeeeeeeee and operates at inference time. (5) MPGD (He al. 2024). This method optimizes the predicted clean data Ty during inference based on a manifold hypothesis. While such refinement incurs additional computational costs during inference, does not fine-tune diffusion model weights, remaining training-fres approach. For more detailed information on each experiment, such dataset, architecture, and baselines, please refer to Appendix E
It is a decoding-basee
Experimental setup. For image experiments (Section 6.1 Section 6.2). we use Stable Diffusion v1.5 (Rombach et al. [2022) as the pre-trained model p (x] c). here e a text prompt (e-g. "cat" or "dog") ") and v is the corresponding image. For the additional control y, we validate compressibilities and aesthetic scores.
6.1 IMAGE: CONDITIONAL GENERATION ON COMPRESSIBILITY
We start by conditioning generations on their file sizes, specifically focusing compressibility Denoting compressibility as CP, we define 4 compressibility labels as follows: Y = 0: CP
0
'Implementing the Classifier-Free haseline in our setting would require using the pre-trained diffusion model to augment x on certain c. Please refer to Appendix E. for details.
Unlike standard tasks in classffer/reconstraciion guidance (Chung et al.) 2022), this score is non- differentiable w.r.t. images. This score function is only dependent on the image itself.
8

---

Published as a conference paper at ICLR 2025
Accuncy
Matro Fl sore
Chosileer- Fice
BRISQUE 4
CLIP Soxe
Recustraction Guidance
MPGD
SMC
0.27 0.0G
0,70 0.0c
0.78
0.22
D.2
Recunatration Guidance
CTRL Chadile--ree
23.2 14.5
10.2 10.1
J05
8---02 35 1.0
228-06
(c) Image quality.
(b) Evaluation.
Labe 20
(a) Histogram.
Carl
Deg
Horee
Rab
-40
Y=0
Compriscibiity
-60
-80
Y=1
-100
Y=2
-140
-160
1500
3000
4500
6000
Reward Queries
Y+3
-
Y=0
Y=1
Y-2
Y-3
(d) Training curves.
(c) Images generated by CTRL.
Figure 1: Results for conditioning on compressibility. Figure plots the histogram of samples generated by the pre-trained diffusion model. Figure T shows the mean compressibility curves during fine-tuning with four distinct lines representing each condition. It is evident that effectivly aligns the generated samples with their target compressibility levels via fine-tuning. b provide evaluation metries, and Figure shows images generated by single model fine-tuned with CTRL.
110.0; == 1: 1110.0 CP 85.0; Y 2: -8.00 S CP < -6000; Y 3 CP -0.00.
Particularly, as depicted in Figure Ta generating samples conditioned on Y = 3 is challenging due to the infrequent occurrence such samples from the pre-trained model.
Results. We evaluate performance across four compressibility levels using the following steps: (1) generating samples conditioned on each Y E [0, 1,2,3] 3): (2) verifying alignment between the generated samples and their conditions; and (3) calculating classification accuracy and macro F1 score. Table Tb presents evaluation statistics. Our results show that CTRL accurately generates samples for each condition, including the rare Y case from the pre-trained model (see Figure Ta). notably outperforming the baselines. Figure I lhooocesss diverse images with correct compressibility levels for various prompts. Additional visualizations are available in Appendix E.5
In Table Ic we provide additional evaluation metrics for baseline methods and CTRL, specifically BRISQUE Mittal et al. 2012) (lower values indicate better image quality) and CLIPScore (Zheng wentai 2023) (higher values reflect better text-image alignment). It is evident that CTRL achieves
better image quality while achieving the highest alignment score.
6.2 IMAGE: MULTI-TASK CONDITIONAL GENERATION
more challenging setting involves multi-tash, conditional generation. this experiment, in addition to compressibility, we simultaneously aim to condition the generations their aesthetic pleasingness. Following prior research (Black et al. 2023 Fan et al., 22033] Uehara et al., (024).. we employ an aesthetic scorer implemented as linear MLP on top of the CLIP embeddings (Radford et al. 2021
which is trained on more than 400k human evaluations.
In this experiment, by leveraging conditional independence of Y1 and Y1 given X (see Example 2), we aim to fine-tune the diffusion model to generate samples with compositional conditions. Specifically, denoting compressibility
CP and aesthetic score AS, we define four compositional conditions
as follows: Y = 0: AS < 5.7, CP < -70; Y = 1 AS < 5.7, CP 2 -70; Y 2
AS 2 5.7, CP < -70; Y 3 AS 5.7, CP > 70. Particularly, depicted in FFiure 2a

---

Published as a conference paper at ICLR 2025
Monkey
Rabbit
Butterfly
Y=0
(a) Histogram.
Y=1
Accursey
Classllie-FFree
Y=2
MPGD
Y=3
0.02
0.02
(c) Images generated by CTRL.
(b) Evaluation.
Figure 2: Results for multi-task conditioning. Figure plots the histogram of samples generated by the pre-trained diffusion model. Table presents the evaluation statistics. Figure displays images generated by single model fine-tuned with CTRL.
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 evaluation procedure outlined in Section 6.1 with results summarized in Table 2b demonstrating that CTRL outperforms all baselines across both tasks by a big margin. Notably, CTRL generate samples rarely produced by the pre-trained model with over 90% accuracy, particularly for the desired class Y = 3 (highly aesthetic images with high compressibility). Producing such images challenging as aesthetically pleasing images typically require more storage and thus have low compressibility. Generated images are displayed in Figure 2c are provided in Appendix E.5
More visualizations
CONCLUSION
We introduce provable RL-based fine-tuning approach for conditioning pre-trained diffusion models on additional controls. Compared to classifier-free guidance, proposed method uses the offline dataset more efficiently and is able to leverage the conditional independence assumption, thereby greatly simplifying the construction of the offline dataset. Our approach is empirically validated across three settings: image generation conditioned on a new task, image generation conditioned on the composition of two tasks, and biological sequence design.
10

---

Published as a conference paper at ICLR 2025
Reproducibility Statement. We submit the code for image experiments supplementary materials. Complete proofs of our theoretical results are provided Appendix DJ Detailed in- formation about experiments, including dataset descriptions, model architecture, and baseline implementations, can be found in Appendix E
REFERENCES
Anderson, B. D. (1982). Reverse-time diffusion equation models. Stochastic Processes and their Applications 12(3), 313-326.
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 Patlern Recognition, pp. 843-852.
Black, K., M. Janner, Y. Du, I. Kostrikov, and S. Levine (2023). Training diffusion models with reinforcement learning.
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, T., B. Xu, C. Zhang, and C. Guestrin (2016). Training deep nets with sublinear memory cost. arXiv preprint arXiv:1604.0617
Chung, H. J. Kim, M. T. Mccann, M. L. Klasky, and C. Ye (2022). Diffusion posterior sampling for general noisy inverse problems. arXiv preprint arXiv:2209.1468)
Clark, K., P. Vicol, K. Swersky, and D. J. Fleet (2023). Directly fine-tuning diffusion models on differentiable rewards. arXiv preprint arXiv:2309.11 17400.
Denker, A., F. Vargas, S. Padhy, K. Didi, Mathis, R. Barbano, V. Dutordoir, Mathieu, U. Komorowska, and P. Lio (2025). Deft: Efficient fine-tuning of diffusion models by learning the generalised h-trrasffrmm Advances in Neural Information Processing Systems 37, 19636-19682.
Dhariwal, and A. Nichol (2021). Diffusion models beat gans on image synthesis. Advances in neural information processing systems 34, 8780-8794.
Epstein, D., A. Jabri, B. Poole, A. Efros, and A. Holynski (2023). Diffusion self- sefd-guidnce for controllable image generation. Advances Neural Information Processing Systems 36, 16222- 16239.
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./6386
Finzi, M. A., A. Boral, A. G. Wilson, F. Sha, and L. Zepeda-Ndñez (2023). User-defined event sampling and uncertainty quantification diffusion models for physical dynamical systems. In International Conference Machine Learning, pp. 10136-10152. PMLR.
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, A., R. Munos, Danihelka, M. Lanctot, and A. Graves (2016). Memory- backprop- agation through time. Advances in neural information processing systems 29.
Guo, C., G. Pleiss, Y. Sun, and K. Q. Weinberger (2017). On calibration of modern neural networks. In International conference machine learning, pp. 1321-1330. PMLR.
Guo, Y., H. Yuan, Y. Yang, M. Chen, and M. Wang (2024). Gradient guidance for diffusion models:
An optimization perspective. arXiv preprint 24X4
Han, X., Zheng, and M. Zhou (2022). Card: Classification and regression diffusion models.
Advances in Neural Information Processing Systems 35, 181 18100-18115.

---

Published as a conference paper at ICLR 2025
He, Y., N. Murata, C.-H. Lai, Y. Takida, T. Uesaka, D. Kim, W.-H. Liao, Y. Mitsufuji, J. Z. Kolter, R. Salakhutdinov, and S. Ermon (2024). Manifold preserving guided diffusion. In The Twelfth International Conference Learning Representations.
Ho, J., A. Jain, and P. Abbeel (2020). Denoising diffusion probabilistic models, Advances in neural information processing systems 33, 1840-68511
Ho, J. and T. Salimans (2022). Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12550
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, 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.0900 09685.
Huang, K., T. Fu, W. Gao, Y. Zhao, Y. H. Roohani, J. Leskovee, 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, K., H. Liu, M. Ryu, O. Watkins, Y. Du, C. Boutilier, Abbeel, M. Ghavamzadeh, and S. S. Gu (2023). Aligning text-to-image models using human feedback. arXiv preprini arXiv:2302. 12192.
Li, X., Y. Zhao, C. Wang, G. Scalia, G. Eraslan, S. Nair, T. Biancalani, S. Ji, A. Regev, S. Levine, et al. (2024). Derivative-free guidance in continuous and discrete diffusion models with soft value-based decoding. arXiv preprini arXiv:2408.08225.
Li, Y. and M. van der Schaar (2023). On error propagation of diffusion models. In The Twelfth International Conference on Learning Representations.
Loshchilov, I. and F. Hutter (2019). Decoupled weight decay regularization. In International Conference on Learning Representations.
Mittal, A., K. Moorthy, and A. C. Bovik (2012). No-reference image quality assessment in the spatial domain. IEEE Transactions on image processing 27(12), 4695-4708.
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 pullern recognition, pp. 2408-2415. IEER.
Nichol, A., P. Dhariwal, A. Ramesh, Shyam, Mishkin, B. McGrew, Sutskever, and M. Chen (2021). Glide: Towards photorealistic image generation and editing with text-guided diffusion models. arXiv preprint arXiv:2 12..041.
Phillips, A., H.-D. Dau, M. J. Hutchinson, V. De Bortoli, G. Deligiannidis, and A. Doucet (2024).
Particle denoising diffusion sampler. arXiv preprint 20l::2202.06.200.
Prabhudesai, M., Goyal, D. Pathak, and K. Fragkiadaki (2023). Aligning text-to-image diffusion models with reward backpropagation. arXiv preprint arXiv:2310.03739
Radford, A., J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, P. Mishkin, Clark, G. Krueger, and I. Sutskever (2021). Learning transferable visual models from natural language supervision. arXiv preprint urXiv:2103.0000
Ramesh, A., M. Pavlov, G. Goh, S. Gray, C. Voss, A. Radford, M. Chen, and Sutskeve (2021). Zero- shot text-to-image generation. In International conference on machine learning, pp. 8821-8831. Pmlr.
Rogers, C. G. and D. Williams (2000). Diffussions, Markov processes and martingales: Volume ltô calculus, Volume 2. Cambridge university press.
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-10000 10695.
12

---

Published as a conference paper at ICLR 2025
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.
Schulman, J., F. Wolski, P. Dhariwal, A. Radford, and O. Klimov (2017). Proximal policy optimization algorithms. arXiv preprint arXiv: 170700777..
Shreve, S. F. et al. (2004). Stochastic calculus for finance II: Continuous-time models, Volume 11. Springer.
Dock-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. J., C. Meng, and S. Ermon (2020). Denoising diffusion implicit models. arXiv preprint arXiv:2010.02550..
Song, J. A. Vahdat, M. Mardani, and J. Kautz (2022). Pseudoinverse-guide diffusion models for inverse problems. In International Conference on Learning Representations.
Song, Y.. J. Dickstiin,,,,, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole (2020). Score- generative modeling through stochastic differential equations. arXiv preprint arXiv:201 13456.
Uchara, M., Y. Zhao, T. Biancalani, and S. Levine (2024). Understanding reinforcement learning- based tee-tuning of diffusion models: A tutorial and review. arXiv preprint arXiv:2407.13744
Uehara, M., Y. Zhao, K. Black, F. Hajiramezanali, G. Scalia, L. Diamant, A. M. Tseng, Bian- calani, and S. Levine (2024). Fine-tuning of continuous-time diffusion models as entropy- regularized control. arXiv preprint 24X1v:2402.15144
Uehara, M., Y. Zhao, K. Black, E. Hajiramezanali, G. Scalia, N. L. Diamant, A. M. Tseng. S. Levine, and T. Biancalani (2024, 21-27 Jul). Feedback efficient online fine-tuning of diffusion models. In R. Salakhutdinov, Z. Kolter, K. Heller, A. Weller, N. Oliver, J. Scarlett, and Berkenkamp (Eds.), Proceedings of the 41st International Conference Machine Learning, Volume 235 of Proceedings of Machine Learning Research, pp. 48892-48918, PMLR.
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. In The Thirty-eighte Annual Conference on Neural Information Processing Systems.
Wainwright, M. J. (2019). Hligh-dimensional statistics: A non-asymptotic viewpoint, Volume 48. Cambridge university press.
Wu, L., B. Trippe, C. Naesseth, D. Blei, and J. P. Cunningham (2024). Practical and asymptotically exact conditional sampling in diffusion models. Advances in Neural Information Processing Systems 36.
Wu, X., K. Sun, Zhu, Zhao, and H. Li (2023). Better aligning eext-to-image models with human preference. arXiv preprint arXiv:2303. 14420.
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-efficien fine-tuning. In Proceedings of the IEEE/CVF International Conference on Computer Vision. pp. 4230-4233
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 arX1v:2304.05977
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, 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.
13

---

Published as a conference paper at ICLR 2025
Zhang, Z., L. Han, A. Ghosh, D. N. Metaxas, and J. Ren (2023). Sine: Single image editing with text-to-ime image diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 6027-6037.
Zhao, S., D. Chen, Y.-C. Chen, J. Bao, S. Hao, 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.
Zhengwentai, S. (2023, March). ccirescore: CLIP Score for PyTorch, https::/github... com/ saited/clip score Version 0.1.0.
14

---

Published as a conference paper at ICLR 2025
A TRAINING DIFFUSION MODELS
In standard diffusion models, given a training dataset {f-1)} ~ P.data(), the goal is to construct
transport that maps noise distribution and data distribution Pdaca E ((x) (X R2d. More specifically, suppose that we have an S
dx, f(t, 0)et a ((t)dwr:
(8)
where f: [0, x TR Rdd is drift coefficient, a [0, 7) - is a diffusion coefficient, wy is d-dimensional Brownian motion, and initial state initial distribution. By denoting the marginal distribution at time by pz.) standard goal in training diffusion models is to learn the parameter 0 so that p (x) Rs Pdata- This means we can
Pini where Piùi E (()) denotes the
(approximately) sample from Pdata by following the SDE ( from 0 to T.
To train diffusion models, we first introduce (fixed) forward reference SDE, which gradually adds noise to Pdata
dz, f(t, 21 )t + o())dw,,
Pdata:
where f: [0, T|R" Rd is drift coefficient, Ø [0. T is the classical denoising diffusion model (Ho et al. 2020). also known as the variance--reeserinng (VP) process, which sets f -0.5z, 1.
+ R is diffusion coefficient. An example
Now, we consider the time-reversal SDE (Anderson 1982). which reverses the direction of SDE
while keeping the marginal distribution, as follows:
dx, f(T t, 2x) + V (og))- di + 5(T t)dw,, To NO la).
(10)
Here, (-) denotes the marginal distribution time for the distribution induced by the reference SDE, and V log (r-<(x;) means derivative w.r.t. Ty, which is often referred to as the score function. Furthremore, when the time horizon I' is sufficiently large, 2, follows Gaussian noise distribution (0, Hence, if we could learn the score function, by following the SDE (TO) 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, natural parameterization is:
f(t, (1; - F(T t, S(T - +)
att) ( -t),
where - x,; 0) is the parametrized neural network introduced to approximate the score function V log (x.) Here, we can leverage the analytical form of the conditional distribution - (+--) (which is Gaussian distribution derived from the reference SDE). This approach enables us to tackle the approximation problem via regression:
i = argmin Exejo, J. [A(0) [|(,, 21;0) V log 1og (11)
where A: [0, T] - IR is a weighting function.
II standard diffuson models, the direction is reversed, i.e., or corresponds the noise distribution.
15

---

Published as a conference paper at ICLR 2025
B PLANNING ALGORITHM FOR CTRL
Algorithm 2 Direct propppropagttion for conditioning
1: Input Batchsize Learning rate '/, Discretization step A,, Exploratory distribution IT
ACC ))-
2: Itinialize: 1 [[0Pre]T,0)
3: for i + 1 to do
4: We obtain n trajectories
(x X (K) in
1k=11 (2) (k)
Z Ik=1-
following (C(#), Y (4)) ~ II(.), ( ~ NO, La), 2o- - 0, and Al
x²a0 ( + g(l 1, ((h,, y(u), x [1) (1) A1 + +o (W)( Awr ~ NOO. (A())),
z = 2 + lg 1, O (()), X v.)) (rr 1, 1(0 1,000,,,111 x
At.
20 (t- - 1)
5: Update a parameter:
Σ k=1
y log P( (x/xxx c(l) - 2 N0
1++1 = =b + V/
6: end for
7: Output: Parameter vis
Inspired by (Clark et al. 2023 Prabhudesai et al. 2023). our planning algorithm, listed in Algorithm 2
is based direct back- proek--rppggtiion This method iterative in nature. During each iteration, (1) compute the expectation over trajectories (E (Ex...7779911))) using discretization techniques such Euler-Maruyama; (2) directly optimize the KL-regularized objective function with respect to parameters of the augmented model (i.e., 4b).
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 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., (0166) to conserve memory, and (c) randomly truncating gradient proo-ppooag to avoid computing through all diffusion steps (Clark et al. [023 Prabhudesai et al., (023).
Remark 4 (PPO). In IllAogriihmm 1 employ direct back pock-propagation (i.e., (i.e.. Algorithmm 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 sugges! using Proximal Policy Optimization (PPO) for planning,
such as Schulman et al. (2017): Black et al. 1) Fan el al. 2023). Other parts remain unchanged.
C INFERENCE TECHNIQUE IN CLASSFIER-FREE GUIDANCE
Although the fine-tuning process sets the guidance level for the additional conditioning (i.e., y) at specific y, classifier-fret free guidance makes it possible to adjust the guidance strength freely during inference. Recall that the augmented model constructed as: g(t, c, y, x;; e)) where =10" Suppose we have obtained drift term ÿ. parametrized by = 10 o J from running Algorithm In inference, we may alter the guidance levels by using the following drift term in the SDE (3)
Uminlt,,,,yu = g(1, 0, 0, xns G) + 71 (,((,,,,,,;;; 0, -0)(
01,, +)) + (2(g(1, e, -) g(t, 0, 2z; 3))
Term 1: pre trained diffusion model conditioned on
Term 2: additional conditioning on y
16

---

Published as a conference paper at ICLR 2025
where indicates the unconditional on or C. In the above, both "/ and y22 do not necessarily need to equal y. They can be adjusted respectively to reflect guidance strength levels for two conditions.
D PROOFS
D.1 IMPORTANT LEMMAS
We first introduce several important lemmas prove main statement.
First, recall that IP (-|c, y) is the induced distribution by the SDE:
dx; g(t, c, y, dt) + (1)dwr: æ0 = Zini
over K conditioning on e and y. Similarly, denote pere (-Ic) by the induced distribution by the SDE:
dxa f c, ze ott) Tu Eni
over K. conditioning on Lemma (KL-constraiined reward). The objective function in 5 is equivalent to
obj B(c.g)-.,,P((ee log [ e) - KL(PP (-1c, v) y) (lPro)(())))))
(12)
Proof. We calculate the KL divergence of 199 and ppre as below
() (-lc) Exa.r Po (11.a) L.'
T (lg u. mz) pree (t, dt
(13)
2
o2(()
This is because
KT((PP(-Ic, y) (Pcc)
dLPS (-I. y) dipree
= Epo(-le.
(lg(t,c,,, xz) - (t,) dt + L (,(t, - -p e, dww
= -Fa(-lea))
2
o²(
(Girsanov theorem)
/' (lg(t,
ftr -
dt
(Martingale property of Itō integral)
(2(1)
Therefore, the objective function in is equivalent
obj E ((eg))
po (ylan, (L(())
(14)
Optimal value function. For the RL problem (5. is beneficial to introduce the optimal optimal value function of (x/ee, y) at any time i E [0, given conditioned on parameters and defined as:
{"
f| (s, x.) C(s.C.y. ds It = x, c,y 02(s)
2f (z|c, max y logp (yler, c)
logp'(ylrr,
2 2
(15)
Specifically, we note that (z(c,c, v) = po (ylæ, c) represents the terminal reward function (i.e., loglikelihood in MDP), while to represents the original objective function (5) that integrates the entire trajectory's 's KL divergence along with the terminal reward.
Below derive the optimal value function in analytical form.
17

---

Published as a conference paper at ICLR 2025
Lemma 4 (Feynman Kac Formulation). At any time t [), given 2; and conditioned and y, we have the optimal value function of (z/c,, y) (induced by the optimal drift term g²) as follows
exp (o] (z/c, y)) (-le [(' (yler, e))² n, c
Proof. From the Hamilton Jacobbi Bellman (HJB) equation, we have
(t)
(zle,)) dr g (ef(, + duž (2)c,,))) dt
l1g f
max
0.
(16)
2
20
where 2(0) is a i-th element in æ. Hence, by simple algebra, we can prove that the optimal drift term satisfies
g" (t,c, = (t, ct, (
By plugging the above into the HJB equation 16 we get
(t)
Pettc pere. Voj dup (x) 2 dr.
o² Vo) 0, (17)
dt
2
which characterizes the optimal value function. Now, using we can show
(1)
d² exx((u (xle,y) + free. V (x)
+ d dexp(u7 (x)cc v))
dr.ddrl
dt
(t)()) 2
Pot(2ll11) drli + fore (ofc, (x)c,)))))
(xlc v)
exp (07 (xlc,,))
Σ
dt
2
=0
Therefore, to summarize, have
(1) d² exp(x; frre, V (exp(e v)) + (xlp(
(18)
2 dabldal
dt
(zlc, v) logp' c)) Finally, by invoking the Feynman Kae formula (Shreve et al. 2004), we obtain the conclusion: Feynman-Kac
(19)
exp (z|c, v)) = Eppet (x c))² F = x, c].
D.2 PROOF OF THEOREM
Firstly, we aim show that the optimal conditional distribution over K. on and y (i.e., IP9" (rlc, v)) equivalent to
ppre (Fle)(p()le)))))) C(e.y) == (xp( v))). C(c,
To do that, we need to check that the above is valid distribution first. This is indeed valid because the above is decomposed into
(ylaT, e))²
ppre (xTlc)
pPre (r|C,
ppre
(20)
C(c,y)
(a2)
(al)
and both (al), (a2) are valid distributions. Especially, for the term (al), we observe
C(c,y) y) = (p" (2)22,, e))² (xrlc)) Eg (ylp)yy/rrr. e)) = exp(vū(2olc,v))
(cf. Lemma 4
18

---

Published as a conference paper at ICLR 2025
Now, after checking 20
is valid distribution, we calculate the KL divergence:
ppre (rlc) (vl
KL
(r)e,
C(c,y)
KL(P" ()( Epo (laa) l (ulor, log C(e,
{%/00
1 (t, c, (.c.,,,.. (t, c, dl y log (yleT, c) + log C(c,
= -Epr (x2))
2
d²(l)
(cf. KL divergence
= v) log C(e,y).
(Definition of optimal value function)
Therefore.
pere (rlc)(
KT. po (rlc.v)
=-66(m0lc log C(c. y) 0.
C(c,y)
Hence,
ppne (P (y KT, c))
1p" (T|c.
C(c.v)
Marginal distribution at t. Finally, consider the marginal distribution at By marginalizing before t, we get
ppte (Tumle)) (p(gg)er,))7//((cc))
Next, by marginalizing after 4,
pire (xle)/C(c,v) x (pp
x, c).
Using Feynman Kac formulation in Lemma 4 this is equivalent to
TP (x/e) exp(u,(rle (llee))))/C v))/C(()))))))
Marginal distribution at T.
We marginalize before T. We have the following pr (xlc)()) (1)
D.3 PROOF OF LEMMA 2
Recall g' (1,c,y,x) frre e, 2) + +02(t) Vz (x/c, y) from the proof of Lemma we have
g'(1,C,,,4) = (r 0) + 2 2() x V log Eppce (-10) (( (ylar,, e)) = xr.c].
E DETAILS OF EXPERIMENTS
E. IMPLEMENTATION OF CLASSIFIER-FRES BASELINE
The effectiveness classifier-free guidance often relies on a sufficiently large offline dataset (c, x y). However, in our experiments (Section 6. and [Section 6.2), only have access to offline datasets (x, y) and (x. /1, y2) respectively. Thus, to implement a classifier-F free guidance baseline in this con- text, leverage the pre-trained diffusion model for data augmentation. The procedure for for[Section 6. is outlined below. The procedure for (Section 6.2 is similar.
Data augmentation. Consider the scenario where Y C and we have access to pere and offline dataset {(I, v. First, use the offline data {xx,)) to train a classifier p X' + A()). We then use P to generate triplets (æ, c, y) ~ px' (2)c)))))(())) for given c. In practice, for text-to-image diffusion models, can be uniformly sampled from set of prompts, such as animals. However, this process becomes computationally demanding when applied to large pre-trained models like Stable
Rombach et al. 2022
19

---

Published as a conference paper at ICLR 2025
Potential limitations. Given the data augmentation strategy, several limitations arise for classifier- free guidance. A primary concern is the accuracy of the trained classifier p(y[z). If the classifier is not sufficiently accurate, the generated values may be unreliable, compromising the quality of the augmented triplets (xt c, y). Additionally, selecting the condition e presents challenges. While models like Stable Diffusion (Rombach al. (022) are pre-trained on vast and diverse datasets with a wide range of prompts. we are constrained to a smaller, more limited set of prompts for in this context. This lack of diversity reduces the representativeness of the augmented data and may lead to mode collapse during fine-tuning common issue observed in fine-tuning of diffusion models (Uehara et al. 2024).
E.2 IMPLEMENTATION OF RECONSTRUCTION GUIDANCE BASELINE
As reviewed in Section 3.2.2
reconstruction guidance baseline employs the following approximation
plyl
p" (ylær, c)p(zr) ar c)dxr = ply) c).c).
where (xt, c) is the (expected) denoised sample given Zt, C, i.c., ( c) [ c]. We note that such approximation is often readily available from diffusion noise schedulers, such DDIM scheduler (Song et al. 2020).
Given such an approximation, we only need to learn py (y pr c) from offline data. Accordingly, we can leverage the trained classifiers from Algorithm (see Seectonn 4.2 Step 2).
As an inference-time technique, the choice of guidance strength is often subtle. We present ablation studies on guidance strength in Appendix E.4 For reporting classification metrics, as shown in Table 16 and Table 2b. we consistently select the optimal configuration for each baseline.
E.3 IMAGES
In this subsection, we provide details of experiments in Section 6 We first explain the training details and list hyperparameters in Table
We use 4 A100 GPUs for all the image tasks. We use the AdamW optimizer (Loshchilov and Hutter 2019) with $1 = 0.9, 32 = 0.999 and weight decay of 0.1. To ensure consistency with also employ training prompts that are uniformly sampled from 50 common animals (Black et al. 2023 Prabhudesai et al. 2023).
previous research, in fine-tuning,
Table 2: Training hyperparameters.
Hyperparameter
compressibility Section 6.1
multi-task Section 6.2
Classifier-frce guidance weight on prompts (i.c., c) " (i.c., strength of the additional guidance on y/)
7.5
7.5
10
10
DDIM steps
50
Truncated back-propagation step
K Uniform(0,500
Uniform(0, 50)
Learning for LoRA modules
le
3c
Learning for the linear embeddings
le 2
1e-2 2
Batch size (per gradient update)
256
512
Number of gradient updates per epoch
2
2
Epochs
15
60
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 (I) 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 [}]] + 1 class labels to embeddings in Rd, where is the same dimension as intermediate diffusion states. Among all embeddings, the first [)]] embeddings correspond to |)] conditions of our interest, whereas the last one represents the unconditional category (i.e., NULL conditioning)
(2) for any E ), 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 [}] embeddings, and freeze the last at zero it is the unconditional label.
20

---

Published as a conference paper at ICLR 2025
We note that, while it is possible to add additional conditioning by reconstructing the networks like ControlNet (Zhang et al. 2023), in practice 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. is especially important to leverage pre-trained diffusion models in our setting where the offline dataset is limited, therefore 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. (202)); Prabhudesai al. (2023): (1) Fine-tuning low-rank adapter (LoRA) modules (Hu et 20021) 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 3: Architecture of compressibility classifier.
Input Dimension
Output Dimension
Layer
x W
ResidualBlock (Conv2d(3, 64, 3x3). BN, ReLU)
64 x
II
W
128 x
4
x W
ResidualBlock (Conv2d(64, 128, 3x3), BN, ReLU) ResidualBlock (Conv2d( (22, 256, 3x3), BN, ReLU)
128
x
256 x
8
x 4
W
x d
256 x
AdaptiveAvgPool2d (1. 1)
256
256
Flatten
256
num cllasses
Linear
Table 4: Architecture of aesthetic score classifier.
Layer
Input Dimension
Output Dimension
Linear
768
1024
Dropout
3
Linear
1024
128
Dropout
5
Linear
128
64
Dropout
7
Linear
64
16
Linear
16
num-classes
Classifiers. In our experiments, we leverage conditional independence for both compressibility and aesthetic scores tasks. Therefore, we only demand data samples (I,, 313) in order to approximate conditional classifier p(yla,, (r, 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 10k images of the AVA dataset (Mur) ray 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 3
aesthetic scores: the classifier is implemented as an MI. on top of CLIP embeddings Rad
ford al. (02T). To train the classifier, which includes more than 250k human evaluations. The specific neural network instruction is listed in Table
the full AVA dataset (Murray al. 2012
Note that in training both classifiers, we split the dataset with 80% for training and 20% for validation. After training, use the validation set to perform temperature scaling calibration Guo et al. (0117).
21

---

Published as a conference paper at ICLR 2025
E.4 ADDITIONAL RESULTS OF RECONSTRUCTION GUIDANCE BASELINE
In this subsection, we provide more results of Reconstruction Guidance for conditioning images on compressibility (see SSction 6.1).
In Figure 3 we plot the confusion matrix for samples generated by Reconstruction Guidance. For each condition, 128 samples are generated and are evaluated. We find that this method struggles to generate samples accurately when conditioned on intermediate labels.
Y=0
88
20
19
1
Conditional control
54
57
99
40
Y=2
57
54
17
Tw3
20
4
36
68
Evaluation
Figure Confusion matrix for Reconstruction Guidance.
Table 5: Results of Reconstruction Guidance conditioned on compressibility. Essentially, guidance
level= 0 indicates that the generations unconditional Y.
Conditional control (Y)
Guidance level (?)
Accuracy 1
Mean score
0.43
-110
5
0.64
-157.4
7.5
0.62
-148.2
10
0.52
-156.5
20
0.66
-152.5
50
0.69
-151.6
0
0.45
-110
0.14
-152.0
7.5
0.12
-189.7
10
0.08
-163.1
20
0.06
- 169.5
50
0
-19.0
2
0.13
-110
2
5
0.02
-104.9
2
7.5
0.10
-122.1
2
10
0.12
-111.3
20
0.08
-157.6
2
50
0.08
-173.5
0
-110
0.46
-71.7
3
7.5
0.46
-67.6
3
10
0.53
-65.6
3
20
0.26
-112.4
50
0.32
-121.5
For completeness, below we provide ablation studies on its hyper- parameters. In Table we present the classification statistics for generations across different conditions (y) and guidance levels (y). Recall that the four conditions are defined as follows:
0: CP -110.0,
-110.0 CP 85.0, Y 2 -85.0 CP -60.0. Y 3: CP -60.0.
22

---

Published as a conference paper at ICLR 2025
Our analysis reveals several key insights:
1. For Y and Y = 3, the accuracy of the generations improves the guidance signal strength increases. This indicates clear positive correlation between the guidance level and the accuracy of generation.
2. Conversely, for intermediate Y and Y 2. guidance signals decrease generation accuracy compared to the pre-trained model, suggesting difficulty in maintaining accuracy within these specific compressibilityy intervals. The challenge in generating samples with medium compressibility scores lies in hand-picking the guidance strength. For instance, generating samples conditioned Y 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 observed across different conditions and guidance levels:
For Y O, 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 y 20 and y 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 = 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 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 according to Table We report the evaluation statistics in Table Tb and provide the confusion matrix in Figure
E.5 ADDITIONAL VISCALIZATIONS
We provide more generated samples to illustrate the performances of CTRL in Figure and Figure 5
Dog
Monkey
Rabbit
Butterfly
Y=0
Ya2
Y=3
Figure 4: More images generated by CTRL in the compressibility task.
23

---

Published as a conference paper at ICLR 2025
Cat
Dog
Horse
Monkey
Rabbit
Butterfly
Peacock
Panda
Y=0
Y=1
Y=2
Y=3
Figure 5: More images generated by CTRL the multi-task conditional generation.
F LIMITATIONS AND MITIGATION STRATEGIES
F.1 COMPUTATIONAL COST
Fine-tuning large pre-trained diffusion models is computationally intensive, especially in multi-task settings. For instance, even in simpler cases, such as fine-tuning a model with single reward function, the process can be slow-DDPO (Black et al. (023) requires approximately 60 A100 GPU hours for fine-tuning Stable Diffusion (Rombach et al. 2002).. In contrast, our single-task experiment Section 6.1] which conditions four compressibility labels (a more complex task than optimizing a single reward), is more computationally efficient, requiring only around 20 A 100 GPU hours. This demonstrates that our method is computationally efficient even when handling more challenging conditioning tasks.
Multi-task experiments Section 6.2 are considerably more demanding, requiring around 200 A100 GPU hours. This is due to the added complexity of aligning the model with multiple control signals, often necessitating larger batch sizes and careful balancing between tasks. As result, multi-task fine-tuning requires not only more GPU time but also careful balancing to prevent overfitting to specific tasks.
Below, we discuss promising ways to help reduce the computational burden without sacrificing performance.
Mitigation strategies. To further reduce computational costs, an effective strategy to truncate backpropagation small fixed number of steps, such as 3 or 5. As noted by (Clark al. (023).. truncating the gradient flow in direct backpropagation to fewer than 10 steps not only sigmificantly reduces computational overhead but also improves optimization stability by mitigating gradient explosion. Interestingly, performance begins to degrade when the number of steps exceeds 10, suggesting that shorter truncation steps (even as few as 1) can be more computationally efficient while maintaining or even improving model performance.
Additionally, mixed precision training can be employed to further accelerate training.
F.2 MEMORY COMPLEXITY
As have clarified in Section 4
many types of off-the-shelf RL algorithms be used for planning.
We recommend using direct back pock-propagation (Clark et al. 2023) or PPO (Black et al. 2023
For direct backpropagation, updating single gradient requires ((L) memory, where L is the number of discretizations. To reduce memory usage, our experiments employed techniques such as (a) fine-tuning only LoRA (Hu al. 2021) modules instead of the full diffusion model, (b) applying
24

---

Published as a conference paper at ICLR 2025
gradient checkpointing (Gruslys et al. 2016 Chen al. 2016), and (c) randomly truncating gradient backpropagation. A detailed discussion of these techniques provided in Appendix B
If memory constraints persist, recommend using PPO for planning, it requires only Q(1) memory per gradient update. Employing mixed precision training can also reduce memory usage.
F.3 CHOICE OF GUIDANCE STRENGTH
7
First, note that 1/7 can be interpreted as the KL weight parameter in standard diffusion model fine-tuning works (Fan et al. 2023; Uchara et al., 2024), where selecting an optimal KL weight remains an open problem,
observed in these works, fine-tuning without entropy regularization often leads to over-optimization. Therefore, introducing a KL weight is beneficial as long as it is neither too small nor too large.
In this work, larger " strengthens the guidance signal of the additional control (see Section 4) but can cause the fine-tuned model to deviate more from the pre-trained model, which is also undesired. For both image experiments, we find that values between 5 and 20 are generally effective. Additionally, even if smaller ? is used during fine-tuning, note that it is possible freely adjust (strengthen weaken) the guidance strength during inference. Details can be found in Appenddx C
y 10 (see Table which provides good balance. In practice,
25


Reviews:
[{"review_id": "pBKgL5YHaQ", "reviewer": "Reviewer_FPqT", "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.", "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:\n $$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)$$\nNow, 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?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "+ This paper is well-organized and the presentation is clear. \n\n+ 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.\n\n[1] Kevin Black, Michael Janner, Yilun Du, Ilya Kostrikov, Sergey Levine. Training Diffusion Models with Reinforcement Learning.", "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.\n\n+ 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "vmCyoFjpv3", "reviewer": "Reviewer_ahRC", "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.", "questions": "See weaknesses.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 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.\n2. 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.\n3. In the experimental section, compared to classifier-free guidance and classifier guidance, CTRL demonstrates superior performance in controllability for image 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.\n2. 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].\n\n[1] Adding Conditional Control to Text-to-Image Diffusion Models. ICCV, 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "EbPrCJGWJo", "reviewer": "Reviewer_acmr", "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.", "questions": "In the error analysis section, the paper briefly discusses three sources of error. Could these three types of errors be more precisely quantified?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 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.\n2. 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.\n3. 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.", "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. \n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "vvV3mEuLH7", "reviewer": "Reviewer_Po57", "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.\n\nNote 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.", "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.\n2. Are there comparisons with more advanced conditional generation models?\n3. 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?\n4. If possible, it is recommended to include a citation to the paper proposing the PPO method [1].\n\n[1] Schulman J, Wolski F, Dhariwal P, et al. Proximal policy optimization algorithms[J]. arXiv preprint arXiv:1707.06347, 2017.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper builds a good theoretical framework for RL-based conditional control to diffusion model.\n2.  As far as I can see, the idea of using reinforcement learning to fine-tune the generation conditions of diffusion models is interesting.\n3. The source code is attached in supplementary material, and I believe it is easy to reproduce.", "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.\n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)