Diffusion Fine-Tuning via Reparameterized Policy Gradient of the Soft Q-Function
Authors: Hyeongyu Kang
Jaewoo Lee
Woocheol Shin
Kiyoung Om
Jinkyoo Park
Venue: ICLR 2026 Poster (2026)
Keywords: Diffusion Models
RL Finetuning
Primary area: generative models
TL;DR: SQDF is a KL-regularized RL method that fine-tunes diffusion models via reparameterized soft Q-gradients, achieving high rewards while avoiding reward over-optimization and preserving diversity.

Abstract: Diffusion models excel at generating high-likelihood samples but often require alignment with downstream objectives. Existing fine-tuning methods for diffusion models significantly suffer from reward over-optimization, resulting in high-reward but unnatural samples and degraded diversity. To mitigate over-optimization, we propose Soft Q-based Diffusion Finetuning (SQDF), a novel KL-regularized RL method for diffusion alignment that applies a reparameterized policy gradient of a training-free, differentiable estimation of the soft Q-function. SQDF is further enhanced with three innovations: a discount factor for proper credit assignment in the denoising process, the integration of consistency models to refine Q-function estimates, and the use of an off-policy replay buffer to improve mode coverage and manage the reward-diversity trade-off. Our experiments demonstrate that SQDF achieves superior target rewards while preserving diversity in text-to-image alignment. Furthermore, in online black-box optimization, SQDF attains high sample efficiency while maintaining naturalness and diversity. 
Our code is available at https://github.com/Shin-woocheol/SQDF.

Published as a conference paper at ICLR 2026
DIFFUSION FINE-TUNING VIA REPARAMETERIZED POLICY GRADIENT OF THE SOFT Q-FUNCTION
Hyeongyu Kang's Jaewoo Leel 2 Woocheol Shin
Kiyoung Oml
Jinkyoo Park
'AIIST *MongooseAl 'Omelet
(khg2000v, jaewoo, woofe, se99an, maxxyoo.parx]kkasaae.
ABSTRACT
Diffusion models excel at generating high-likelihood samples but often require alignment with downstream objectives. Existing fine-tuning methods for dif- fusion models significantly suffer from reward over-optimization, resulting high-reward but unnatural samples and degraded diversity. To mitigate over- optimization, we propose Soft Q-based Diffusion Finetuning (SQDF), a novel KL-regularized RL method for diffusion alignment that applies reparameterized policy gradient of training-free, differentiable estimation of the soft Q-function. SQDF is further enhanced with three innovations: discount factor for proper credit assignment in the denoising process, the integration of consistency models refine Q-function estimates, and the use of off-policy replay buffer to im- prove mode coverage and manage the reward-diversity trade-off. Our experiments demonstrate that SQDF achieves superior target rewards while preserving diver- sity in eexttto-mmage alignment. Furthermore, in online black-box optimization, SQDF attains high sample efficiency while maintaining naturalness and diversity. Our code is available at https: thub. com/ Shi n-woocheo //spch SQDF.
INTRODUCTION
Diffusion models (Ho et al., 2020: Song et al., 2021b) have emerged as dominant paradigm for generative tasks, achieving high-fidelity sample generation in domains such as text-to-image synthe- sis (Rombach et al., 2022), video generation (Ho et al., 2022), and biological molecules (Lee et al., 2025). This can framed as an optimization problem: maximizing a reward signal that quantifies desired property, such as aesthetic quality (Schuhmann, 2022), rext-to-image alignment (Xu al., 2023; Wu et al., 2023), or molecular bioactivity (Gosai et al., 2023).
Existing fine-tuning methods for reward optimization are broadly categorized into two main ap- proaches: reinforcement learning (RL)-based methods (Black et al., 2023) and direct backpropaga- tion methods (Xu et al., 2023; Clark et al., 2023; Prabhudesai et al., 2023). Although these methods effectively optimize rewards, their singular focus on reward maximization makes them highly prone over-optimization (Skalse al., 2022), often the expense of sample quality and diversity. While KL-divergence regularization with pre-trained model is proposed to mitigate over-optimization (Ue- hara et al., 2024a), they often necessitate explicit value function training notoriously unstable process (Uchara et al., 2024b; Hu et al., 2025; Zhou et al., 2024) or depend on high-variance Monte Carlo gradient estimators (Fan al., 2023; Venkatraman et al., 2024). Although downstream reward gradients are powerful training signals, directly leveraging them to fine-tune diffusion models without over-optimization remains an open problem.
To this end, we propose Soft Q-based Diffusion Finetuning (SQDF), method that employs reparameterized policy gradient guided by training-free soft Q-function within a KL--regularized RL framework. The of our approach is to approximate the soft O-function via single-step pos- terior mean approximation (Li et al., 2024), a strategy that circumvents the need for unstable value function learning. This approximation is differentiable under the parameterized oracle proxy models (Xu et al., 2023; Wu et al., 2023; Uehara et al., 2024c), enabling the direct use of gradient for low-variance and cample-efficieee policy updates (Haarnoja et al., 2018). The entire approach is
"qqual contribution authors.

---

Published as a conference paper at ICLR 2026
Step1. Sampling from Pe
Step2. Train diffusion model Pe
Diffusion Model
Consistency Model
Reward Model
Replay Buffer
Store
Sample
Po
R
%0
Iteration
Backoward
Forward
Training *
Frezen
Figure 1: Overview of the SQDF framework. The process involves two stages: (1) samples gener- ated by the diffusion model pe are stored in a replay buffer; (2) a noisy sample 2y is drawn from the buffer and denoised one step by the diffusion model pg. The consistency model fe then takes 21-1 as input and predicts the clean sample λn- This prediction is evaluated by a reward model Te. and the resulting reward gradient is used to update po via reparameterized policy gradient.
guided by a KL-regularized objective (Uchara et al., 2024c), which enforces the finetuned model to remain close to the pretrained data distribution, thereby preserving sample quality and diversity.
Beyond the foundation, SQDF introduces three additional innovations. (1) Incorporation of a dis- count factor to downweight early denoising steps, reflecting their limited influence final sample quality (Ho et al., 2020). (2) Integration of consistency models (Song et al., 2023), refining estima- tion quality of soft-Q function. (3) Off- Off-policy updates with replay buffer allow to exploit its benefits for mode-coverage (Sendera et al., 2025), managing the reward-diversity trade-off by curat- ing the data used for the finetuning of the diffusion model.
We evaluate SQDF in two settings. First, we apply SQDF to optimize Stable Diffusion 1.5 (Rombach al., 2022) and Stable Diffusion XL (Podell et al., 2023) using differentiable rewards, specifically LAION aesthetic score (Schuhmann, 2022) and the HPSv2 human preference score (Wu et al., 2023). Our results demonstrate that SQDF effectively optimizes rewards while mitigating over- optimization. Second, we test SQDF in a black-box optimization problem, where optimizing reward and preserving diversity are both crucial under limited query budget (Bengio et al., 2021). this experiment, SQDF achieves high sample efficiency while preserving diversity.
2
RELATED WORKS
2. ENTROPY REGULARIZED REINFORCEMENT LEARNING
Entropy regularized RL augments the expected return with Shannon entropy bonus KL diver- gence penalty. RL with Shannon entropy bonus, usually called as maximum entropy RI, enhances exploration and robustness, enabling deep RL methods achieve strong performance in various domains (O' OO'Donoghue
al., 2017; Haarnoja et al., 2017; 2018). KL-regularized RL encourages high-return actions while constraining the policy to remain close to a reference distribution (Wu et al., 2019; Nair et al., 2020). In SQDF, employ KL-regularized RL framework with soft Q-function to constrain the policy to stay close to the pretrained diffusion model, thereby preserving naturalness and diversity of the pretrained diffusion model (Uchara et al., 2024a).
2.2 ALIGNMENT OF DIFFUSION MODELS FROM HUMAN FEEDBACK
The expressive capability diffusion models in capturing data distributions effective across var- ious domains (Yang et al., 2023), but aligning pre-trained diffusion models with downstream objec- tives is often necessary. Black et al. (2023); Fan et al. (2023) employ PPO (Schulman et al., 2017b) to fine-tune the diffusion model by maximizing the black-box reward function. Other works (Xu et al., 2023; Prabhudesai et al., 2023; Clark et al., 2023) assume the reward gradient accessible, and directly backpropagate gradient signals from the final state of the diffusion process.
2

---

Published as a conference paper at ICLR 2026
While effective at optimizing rewards, they prone to over-optimization. leading semantic collapse and reduced diversity (Skalse et al.. 2022: Gao et al., 2023). To address this issue, KL regularization approaches introduced. While effective in mitigating over-optimization, Uchara et al. (2024b) requires value function training, notoriously hard to train diffusion (Hu et al., 2025), well as initial noise distribution training. Fan al. (2023); Venkatraman et al. (2024) estimate the gradient signal with high-variance Monte Carlo estimation even if there exists a reward gradient al. (2025) introduce stochastic optimal control for
signal available. Recently, Domingo-Enrich finetuning flow-base bwsdaaed models to reward tilted distribution. SQDF leverages training-free, one-step soft Q-function approximation to directly apply the reward gradient via reparameterized policy update, effectively fine-tuning the model without backpropagation through the denoising chain.
3 BACKGROUND
3.1 DIFFUSION MODELS
Diffusion model is generative model that approximates the data distribution, generating samples by gradually denoising the Gaussian noise zr ([rr; 0, I) via the parameterized reverse process (o(rr) The reverse process is defined Markov chain as follows:
po( p(xT) II (o [x)),
pa xre Mar-1 na pa 1), of
The forward process ( (o)), is Markov chain that gradually noises the zo from the data distri- bution according to the fixed variance schedule {3,)z=1-
( 20) II M(z: /1 B, x1-1, 3,1), (< v B, 3e1) (2)
1=1
Diffusion models can be trained with the following noise prediction loss (Ho et al., 2020):
((0) EE somppo t=id((1.....))) t=, (o.)
ed(Vor + v1 ā.e, t) 1)
where āg = II (1 - B.) with noise parameterization eg (Zz, t).
3.2 MARKOV DECISION PROCESS FOR DIFFUSION FINETUNING
The diffusion reverse process can be naturally formulated as a Markov Decision Process (MDP) due to its inherent Markov property. MDP is defined as M = AS, T, R, Po), where S denotes the state space, A the action space, T S x + S the transition function, R. S -> IE the reward function, and Po the initial state distribution. We formulate finite-horizon MDP with sparse rewards and the deterministic transition as follows:
St (IT-I,T- t) Ta(a:)s:) A pa po( P(sc+1l5t, az) 2 112--1111--1111
R(s,, ax) & (r(2)) =
Qr A FT-1-1
Po(su) & (8r,
10 otherwise
3.3 KL-REGULARIZED RL IN DIFFUSION MDP
We define the KL-regularized objective for diffusion model alignment, using the MDP formula- tion in Section 3.2, following Uchara et al. (2024a):
p arg max (x)pu(T) Pe
r(xn)
Σ 1=1
DKI(pa(-
(x))ll' 1=:))lo-- xx))
(4)
where 7 = (Dr, trained diffusion model as reference policy. we define the optimal soft Q-function for a state-action pair (I,, xt-1) the sum of reward with KL regularization:
I1, FD) denotes denoising trajectory generated by po. and p' denotes pre-
De (pppppp"aaaa (+Ixx)llp'( (-(((((((=)))) (00 (aa))
Quaa (1e, 11-1) Ep-
r(x))
(5)
3

---

Published as a conference paper at ICLR 2026
then analytically derive the corresponding optimal policy:
p" (-1x4) ) = arg max Bu.-i-po((le) [Qio (t-1) - aD (=a)ll>))) (++x))
(6)
pe(-)a;)
To obtain the soft optimal policy in Equation (6), we need the optimal soft Q-function which is characterized by the soft Bellman equation:
Quaa(an 2z-1) = ((( + Vion(2i-1))
R(F,, + Vion (1o1(1111))
Vo (x)) = a log E Be--r=(((la))
exp
(8)
After recursively solving the soft Bellman equation, the soft Q-function can be approximated as: (r(z))
Quon (Ji, ft-1) a log B........... a--1) exp
-t-1 ((
where (2te) Ey (rr [are] is the posterior mean approximation driven from Tweedie's for- mula (Efron, 2011; Chung et al., 2023). We provide details in Section A.
4 METHODS
This section details our proposed method, SQDF, a novel KL-regularized RL method for fine-tuning diffusion models while mitigating reward over-optimization. The core of SQDF lies in bridging the signal from the reward model directly to the intermediate denoising process of the diffusion through the waining-frree soft Q-function. To enhance the stability and effectiveness, we introduced three components. (1) discount factor y, (2) consistency model, and (3) experience replay buffer. We illustrate the overview of method in Figure I.
4.1 POLICY IMPROVEMENT VIA REPARAMETERIZA POLICY GRADIENT
Gradients from a differentiable reward function offer a potent signal for fine- sene-tuning diffusion models (Clark et al., 2023; Xu al., 2023), yet existing KI KL-regularized RL methods struggle leverage them directly (Uehara et al., 2024b; Venkatraman et al., 2024). SQDF directly exploits the reward gradient serve as the gradient of the soft Q-function, approximated in training-free manner in Equation (9). This approach eliminates the need to train separate Q-network and avoids its associated instabilities (Zhou et al., 2024; Hu al., 2025).
Leveraging the reward gradient as the Q-function gradient allows us to employ the reparameterized policy gradient, which provides a low-variance and accurate training signal for reward optimiza- tion. By substituting the soft Q-function in Equation (6) with the soft Q-function approximation in Equation (9), we can derive reparameterized policy gradient loss with the reward function:
(( Ex,, E po(la.)) Qun ( (11, x2-1)
(x=-1
(10)
Ex,, [ (-1.)) ( (1 + oD (po (2;-1//11)))) (x))) (x=]]
(11)
Since Jt-1 sampled from Pa (-1x;) a stochastic variable, the gradient from r(āo) cannot be back- propagated. Therefore, by using reparameterization with 2T-1 = ua(2x, + a (Kingma & Welling, 2013), we utilize the gradient signals can utilized to update the policy parameters (Lill- icrap et al., 2015; Haarnoja al., 2018). Then the gradient of the Equation (11) is given by:
C(C = Ez, [EE-N((,)) [-V
Vopo(ze, t) + a aV (poll']]]]
(12)
4.2 STABILIZATION AND OFF-POLICY TECHNIQUES IN SQDF
Although the reparameterized policy gradient in Equation (10) directly leverages the powerful re- ward gradient signal, it still faces significant challenges. In the early denoising steps, due to lower signal-to-noise ratio (Kingma et al., 2021), I), each denoising step has a minimal influence on the final generated samples (Ho et al., 2020). Moreover, approximation error from Equation (9) exacerbates in early denoising steps, which leads to unreliable training siganl. To address these issues, SQDF incorporates two stabilization techniques: the adoption of a discount factor and consistency mod- els. Separately, with the inherent off- off-policy training of Equation (1 1), we utilize buffer D to control reward-diversity trade-offs.

---

Published as a conference paper at ICLR 2026
(a) DDPM
(b) Tweedie's formula
(c) Consistency model
1=
t= 30
1
t=
= t= 20
Figure 2: Comparison multi-step sampling with one-step æ estimation. (a): DDPM 50-step sampling accurately capture Zy distribution. (b): A one-step Ty estimation via Tweedie's formula is highly inaccurate, particularly at early denoising steps. (c): Consistency model, however, provides an 2g estimate with uniform accuracy across all timesteps.
4.2.1 7 DISCOUNTED MDP FOR IMPROVED CREDIT ASSIGNMENT
Due to inherent stochasticity in the diffusion reverse process, the influence of action at 2; on the final sample Tu diminishes as t increases (Ho et al., 2020; Song al., 2021b). Prior MDP formulation (Black et al., 2023) typically sets ? = 1, assigning equal credits to actions taken across all timesteps. For better credit assignment, the credit for an action at timestep by 7², exponentially down-wcighting actions in early denoising steps in large align the optimization with their limited contribution to the final sample. Then we can rewrite the KL-regularized RL objective at Equation (4) with discount factor ry as follows:
introduce discount factor 7 E [0, 1]. By setting 7 1, it scales
p arg max E By-po(T) yT -T
Σ "Dr 2x)lp'( |x)))
(13)
In the case of discounted MDP, a recursive expansion of the soft optimal value function denoted at Equation (8) leads to the following bounds for Q (Zte, 1t-1:
a log Ex exp (t" r((o))] I Qion (2t, xx-1) S r a log Exo. exp
(((o))
(14)
From the above inequalities, we observe that applying posterior mean approximation to both the lower and upper bounds of the soft Q-function yields the identical first-order approximation. Moti- vated by this observation, we approximate the soft Q-function under the discounted MDP as:
Quaa (It, 11) R 7 ((
(15)
The detailed derivation of the above equations can be found in Section B.
4.2.2 CONSISTENCY MODEL FOR BETTER SOFT Q ESTIMATION
The critical limitation of Tweedie's formula for approximating the posterior mean èn is its lack of reliability under high noise levels (Chung et al., 2023). This issue illustrated in Figure 2-(b), where Tweedie's 's formula often yields samples that lie far outside the data distribution with the early diffusion timesteps. Although using a discount factor 7 can mitigate the effects of inaccurate approximation, it can still mislead the policy update process (Jain et al., 2025). Although employ- ing multi-step Ordinary Differential Equation (ODE) sampling (Song al., 2021a) suitable for accurate mean posterior approximation, it is impractical due to the requirement of backpropagation through the denoising trajectory, which leads gradient instability (Clark et al., 2023).
To address these limitations, we employ the consistency model (Song et al., 2023) (are to improve the accuracy of the posterior mean approximation. Consistency models are trained to map the noisy input He to the corresponding clean image xo by distilling integration of probability flow ODE (Lu & Song, 2024). By using the consistency model as a reference diffusion model p', we achieve a more accurate posterior mean than that provided by Tweedie's 's formula. As shown in Figure 2-(c), the consistency model provides a better In approximation. We argue that accurate posterior mean prediction improves the soft-Q function approximation and better guides training((See Section 5.3.)
5

---

Published as a conference paper at ICLR 2026
4.2.3 OFF-POLICY TRAINING WITH REPLAY BUFFER
Unlike previous works that rely on on-policy samples (Black al., 2023; Clark et al., 2023; Comingo-Ennichh et al., 2025), the proposed SQDF loss in Equation (10) naturally accommodates off-policy updates. We integrate replay buffer, denoted as D, along with an experience replay strategy explained in Section E. This approach allows for the reuse of rare, high-reward, and di- verse samples, facilitating balance the reward-diversity trade-off and enhancing mode coverage. Finally, combining all the key factors of the SQDF, we can define SQDF loss as follows:
(0) Ba,DD, me i~po 1-7 1- (Jr (x+-1)) 1)) + DRL. (pa (po ( (x))llp (x [x)]]] (16)
EXPERIMENTS
This section presents comprehensive empirical analysis of SQDE, For all tasks, utilize Stable Diffusion v1.5 (Rombach et al., 2022), and all experiments are conducted using three random seeds. Our experiments aim answer the following research questions:
Does SQDF optimize the reward while mitigating verr-optimization? ($5.1)
Can SQDF effectively finetune diffusion models in online black-box optimization settings? ($5.2)
How do the individual components of SQDF contribute to improved finetuning performance in terms of efficiency, diversity, and stability? ($5.3)
5.1 FINETUNING FOR TEXT-TO-IMAGE DIFFUSION MODELS
For finetuning the text-to-image diffusion models, we employ the LAION aesthetic predic- tor (Schuhmann, 2022) and human preference scores (HPSv2) (Wu et al., 2023) objective func- tions. Additional details of the task and algorithms can be found in Section D.2.2 and Algorithm 1.
5.1.1 EXPERIMENT SETUP
Evaluation metrics: As shown in prior works (Black et al., 2023; Clark et al., 2023), reward over- optimization often leads to unrecognizable semantic content and reduced diversity. We refer to these two phenomena semantic collapse and diversity collapse respectively. To assess semantic col- lapse, we employ the prompt alignment scores such as HPSv2 (Wu et al., 2023), ImageReward (Xu et al., 2023), and PickScore (Kirstain et al., 2023), which are trained to align text-image pairs with human preference datasets. To detect diversity collapse, we employ two diversity measures: the mean pairwise distance computed with LPIPS (Zhang et al., 2018) and the mean pairwise cosine distance calculated using DreamSim features (Fu et al., 2023). We provide visual examples of se- mantic and diversity collapse cases in Section C.
We compare against: (1) RL-based method: DDPO (Black et al., 2023), (2) direct back- propagation method: DRaFT (Clark et al., 2023), ReFL. (Xu al., 2023), and (3) KL--regllarized variants of DDPO and DRaFT, which introduce an auxiliary Kl.-regularization term in the reward. Further details of KL-regularized baselines provided in Section D.
Baselines:
5.1.2 RESULTS
SQDF mitigates over-optimization: Figure shows how alignment and diversity scores change as the target reward is optimized. The first row demonstrates the optimization of the aesthetic score, and the second row illustrates the HPS optimization. Panels (a) show that gradient- based base- lines such as ReFI. and DRaFT achieve high aesthetic scores but suffer sharp declines in alignment and diversity, highlighting their vulnerability to reward over-optimization. DDPO, which does not exploit gradient signals, fails to reach comparable aesthetic score and exhibits rapid diversity col- lapse. For optimizing HPS score, (c)-(h) demonstrate that SQDF consistently achieves the highest alignment and diversity at equivalent reward levels, demonstrating its generalization capabilities.
Comparison with KL-Augmented Baselines: natural question is whether the benefits of SQDF can be replicated by simply adding KL-divergence term existing methods. To address this, compare the trade-off curves generated by sweeping the regularization strength a for both SQDF and other KI -L-augmented baselines. The results, depicted in Figure 4, show that SQDF achieves higher

---

Published as a conference paper at ICLR 2026
(a) Aesthetic vs. Alignment
(b) Aesthetic vs. Aligument
(c) Aesthetic vs. Diversity
(d) Aesthetic vs. Diversity
0st-
ImageRoward
DreamSim- Div
HPS
0.55
Aesthetic (c) HPS vs. Alignment
Aesthetic
Aesthetic
Aesthetic
(1 HPS vs. Alignment
(g) HPS vs. Diversity
(b) HPS vs. Diversity
A22
0.55
ImayeReward
I.P.S--DD DiV
DreanSim Div
0.4
Pick
021
0.
-35-
LIPS
1IPS
IIPS
IIPS
Our
DDPO
DRaFT
ReFL
Figure 3: Comparison of evaluation metrics during optimization of the target reward. Top: The target reward is the LAION aesthetic score. Bottom: The target reward HPSv2. (a), (b), (e), and (f): evaluation of alignment score using ImageReward and HPS. (c), (d), (g), and (h): evaluation of diversity using LPIPS and DreamSim.
(a) Aesthetic vs. Alignment
(b) Aesthetic vs. Alignment
cc) Aesthetic vs. Diversity
(d) Aesthetic vs. Diversity
Image Reward
DreamSim Div
+100
Aesthetic
Acsthetic
Aesthetic
Aesthetic
Ours
DDPO+KL
DRaFT-KL
Figure 4: Comparison of trade-off curves with KL-regularized baselines. Curves are obtained by varying the KL-regularization coefficient a. Darker points correspond to stronger KL-regularizer.
rewards while maintaining better alignment and diversity, mostly occupying the Pareto optimality across different metrics. These findings suggest that SQDF provides a more effective framework for optimizing the KL-regularized objective with reward gradient backpropagation while mitigating over-optimization. Qualitative comparisons can be found in Figure 5 and Section E.
5.2 ONLINE BLACK- BOX OPTIMIZATION WITH DIFFUSION MODELS
Given a limited query budget for evaluating an oracle reward function, online black-box optimiza- tion as suggested by Uehara et al. (2024c) involves fine-tuning a generative model to achieve high oracle reward scores while preserving the model's naturalness and diversity. In this experiment, we adopt the experimental setting from (Uchara et al., 2024c), iteratively fine-tune a diffusion model with aesthetic score considered black-box oracle function. Additional details of the task and algorithms can be found in Section D.2.3 and Algorithm 2.
5.2. EXPERIMENT SETUP
Baselines: We compare SQDF against: (1) SEIKO (Uehara et al., 2024c), KL-regularized direct backpropagation approach that utilizes the reward signal of the proxy reward model. We evaluate two variants: SEIKO-UCB and BoottaBo which introduce uncertainty bonus improve

---

Published as a conference paper at ICLR 2026
Pretrained Reward: 5.09
Ours Reward: 8.22
ReFL Reward: : 7..1
DRaFT Reward 8.08
DR Reward: 6.62
DDPO Reward 6.92
DDPO+KL Reward: 7.14
Figure 5: Comparison of generated images from different fine-tuning methods, using model check- points selected when reward of 8.0 was achieved (or the maximum reward if 8.0 was not reached). The average reward for the presented images is shown for each method.
Method
Target (Aesthetic)
Alignment
Diversity
(ewageReward)
(Hps)
(Lpsis-Div)
(Dreamsim- Div)
PPO+KL
6.63 (0.45)
-1.35 (1.30)
0.24 (0.01)
0.47 (0.07)
0.44 (0.15)
SEIKO- BootttBootstrap
7.80 (0.11)
-1.69 (0.64)
0.23 (0.01)
0.36 (0.08)
0.24 (0.13)
SEIKO-UCB
7.49 (0.18)
-1.08 (0.96)
0.24 (0.01)
0.40 (0.12)
0.32 (0.24)
SQDF-Bootstrap
7.87 (0.11)
1.14 (0.08)
0.27 (0.00)
0.49 (0.02)
0.53 (0.02)
SQDF-UCB
7.87 (0.15)
(0.12)
0.27 (0.00)
0.51 (0.03)
0.54 (0.04)
Table 1: Results of online black-box optimization. SQDF achieves superior performance under the
same oracle query budgets. We report the mean and standard deviation over random seeds.
exploration. (2) PPO + KL, where the policy is fine-tuned directly with oracle rewards via KL- penalized PPO update (Schulman al., 2017b).
Training surrogates and uncertainty bonus: We train our proxy of the reward model and introduce an uncertainty bonus following SEIKO. Detailed in Section D.1.2.
5.2.2 RESULTS
Table reveals clear superiority SQDF across all evaluation metrics. While SQDF consistently achieves the highest target rewards alongside strong alignment and diversity scores, SEIKO, the di- rect propagation method, suffers from a critical trade-off: both variants show deteriorating alignment and reduced diversity optimization progresses. This stark contrast demonstrates that even under an identical KI. regularization framework, our approach fundamentally outperforms existing meth- ods. Unlike controlled fine-tuning tasks, this MBO setting with imperfect reward proxies typically drives models out of distribution, losing the naturalness. Interestingly, SQDF maintains robustness across all objectives simultaneously, proving its robustness to inaccurate rewards. Figure offers visual comparison between SEIKO and SQDE.
5.3 ABLATION STUDY
To provide deeper insights into SQDF's effectiveness, present comprehensive ablation studies examining each of the three key techniques outlined Section 4.2.
SQDF component ablation: First, examine the effeet of the discount factor 7. In the Equa- tion (13), prior works implicitly use 7 1, whereas our setting adopts y E [0. 1). As shown in Figure 6, removing the discount factor (7 ultimately achieves higher aesthetic score. How-
8

---

Published as a conference paper at ICLR 2026
Dream PreamSSim-i Sim- Lin-iiv
Aesthetic
13
ImageRewart
LPIPS-DiS --is-Div
7/ 6.5
53
epoch
epoch
epoch
epoch
Ours
Without discount
Figure 6: The effect of the discount factor training dynamics. Comparison of method against a bascline where the discount factor is set to y = 1.
ever, undiscounted SQDF exhibits slower optimization at early epochs and a significant drop in alignment and diversity scores. This suggests that introducing small credit near the initial state of the diffusion process improves credit assignment and reduces the high approximation error in the early denoising steps. We report the result of varying in Section E.2.
Table presents the performance of SQDF when the consistency model replay buffer is ablated. Remov- ing the consistency model decreases performance on the target reward. This indicates that the use of the con- sistency model improves training efficiency by making the soft Q-function upproximation more reliable, dis- cussed in Section 4.2.2. Removing the buffer reduces diversity scores, as also shown in Section E.1. sug- gesting that the buffer helps preserve model support by reusing past experience. This aligns with prior findings that replay buffers improve mode coverage and mitigate catastrophic forgetting (Mnih et al., 2013).
Evaluation Metric
Method
Aesthetic
DreamSimi-Dr
LPIPS-Div
SQDF
7.87
0,58
0.56
wlo CM
7.10
0.62
w/o buffer
8.06
0.56
0.55
Table Ablation study on Consistency model (CM) and Buffer. The consistency model enables faster convergence, while the buffer preserves diversity.
The naive one-step prediction of ig using Tweedie's s formula is highly inaccurate in the early denoising steps, illustrated in Figure 2- (b). A natural alternative to use n-step ODE sam- pling with DDIM (Song et al., 2021a) for accurate ap- proximation. As shown in Figure 7, SQDF with 2- step DDIM sampling substantially improves optimiza- tion compared to Tweedie's 's formula. However, 4-step DDIM leads unstable training, which we attribute to high variance in the gradients. We adopt the consis- tency model as Pareto solution for reliable za predic- tion and stable training.
DDIM vs consistency:
6.75
SQDF
I-step
6.50
2-step
Aesthetic
4-step
6.25
50
100
150
200
Epoch
Figure 7: DDIM n-step versus consis-
tency model 2o predictor.
CONCLUSION
In this work, we introduce SQDF, novel KL.-regularized reinforcement learning framework for fine-tuning diffusion models that directly utilizes reward gradients. The core of our method is reparameterized policy gradient with training-frre,, one-step approximation of the soft Q-function. This allows SQDF to efficiently optimize for downstream objectives while effectively mitigating reward over-optimization. We further enhance training stability by incorporating three key compo- nents: a discount factor for credit assignment, a consistency model for reliable soft Q estimation, and a replay buffer manage the reward-diversity trade-off. Extensive experiments demonstrate that SQDF successfully optimizes target rewards while preserving naturalness and diversity, thereby pushing the Pareto frontier in both text-to-image fine-tuning and black-box optimization tasks. We believe that leveraging more advanced one-step distillation models and sophisticated buffer manage- ment techniques are promising directions for future work.

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGEMENT
We thank the anonymous reviewers for their insightful comments and suggestions, which signif- icantly improve our manuseript. This work is supported by the National Research Foundation of Korea (NRF) grant funded by the Korea government funded by the Korea government (MSIT) (No. 15-2024-00410082), and NRF grant funded by the Korea governeen((MSTT) (No. RS-2025-00563763)
NRF grant
THE USE OF LARGE LANGUAGE MODELS
Large Language Models were employed exclusively for two auxiliary tasks: (1) minor polishing of the manuscript text for improving grammar and readability, and (2) limited assistance in code implementation for debugging syntax or refactoring functions. Importantly, LLMs did not contribute to the conception of the research problem, the development of the core methodology, or the design and execution of the experiments. All critical ideas, methods, and analyses presented in this paper the original work of the authors.
REFERENCES
Emmanuel Bengio, Moksh Jain, Maksym Korablyov, Doina Precup, and Yoshua Bengio. Flow network based generative models for non-iterative diverse candidate generation. Advances in neural information processing systems, 34:27381-27394 2021.
Kevin Black, Michael Janner, Yilun Du, Ilya Kostrikov, and Sergey Levine. Training diffusion models with reinforcement learning. In The Twelfth International Conference on Learning Rep- resentations, 2023.
Hyungjin Chung, Jeongsol Kim, Michael Thompson Mccann, Mare Louis Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. In The Eleventh Interna- tional Conference on Learning Representations, 2023.
Kevin Clark, Paul Vicol, Kevin Swersky, and David Fleet. Directly fine-tuning diffussion models on differentiable rewards. In The Twelfth International Conference on Learning Representations, 2023.
Carles Domingo- Enrich,, Michal Drozdzal, Brian Karrer, and Ricky TQ Chen. Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control. In The Thirleenth International Conference Learning Representations, 2025.
Bradley Efron. Tweedie's formula and selection bias. Journal of the American Statistical Associa- tion, 106(496):660216611 2011.
Ying Fan, Olivia Watkins, Yuqing Du, Hao Liu, Moonkyung Ryu, Craig Boutilier, Pieter Abbeel, Mohammad Ghavamzadeh, Kangwook Lee, and Kimin Lee. Dpok: Reinforcement learning for fine-tuning text-to-image diffusion models. Advances Neural Information Processing Systems, 36:79858 79885, 2023.
Stephanie Fu, Netanel Y Tamir, Shobhita Sundaram, Lucy Chai, Richard Zhang, Tali Dekel, and Phillip Isola. Dreamsim: learning new dimensions of human visual similarity using synthetic data. In Proceedings of the 37th International Conference on Neural Information Processing Systems, pp. 50742-50768, 2023.
Leo Gao, John Schulman, and Jacob Hilton. Scaling laws for reward model overoptimization. In International Conference on Machine Learning, pp. 10835-10866. PMLR, 2023.
Sager Gosai, Rodrigo Castro, Natalia Fuentes, John C Butts, Susan Kales, Ramil R Noche, Kousuke Mouri, Pardis C Sabeti, Steven K Reilly, and Ryan Tewhey. Machine-guided design of synthetic cell type-specific cis cis-regulatory elements. bioRxiv, 2023.
Tuomas Haarojja, Haoran Tang. Pieter Abbeel, and Sergey Levine. Reinforcement learning with deep energy-based policies. In International conference on machine learning, pp. 1352-1361. PMLR, 2017.
10

---

Published as a conference paper at ICLR 2026
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International confer- on machine learning, pp. 1861-1870. Pmlr, 2018.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851, 2020.
Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David Fleet. Video diffusion models. Advances in neural information processing systems, 35:8633- 8646, 2022.
Edward J Hu, yelong shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation of large language models. In International Con- ference on Learning Representations, 2022.
Zijing Hu, Fengda Zhang, Long Chen, Kun Kuang, Jiahui Li, Kaifeng Gao, Jun Xiao, Xin Wang. and Wenwu Zhu. Towards better alignment: Training diffusion models with reinforcement learning against sparse rewards. In Proceedings of the Computer Vision and Pattern Recognition Confer- ence, pp. 23604-23614, 2025.
Vineet Jain, Kusha Sareen, Mohammad Pedramfar, and Siamak Ravanbakhsh. Diffusion tree sam- pling: Scalable inference-time aligmment of diffusion models. In Second Workshop Test-Time Adaptation: Putting Updates to the Test! at ICMI. 2025, 2025.
Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Ad- vances neural information processing systems, 34:21696-21707 2021.
Diederik Kingma and Max Welling. arXiv:1312.6114 14, 2013.
Auto-encoding variational bayes. arXiv preprint
Yuval Kirstain, Adam Polyak, Uriel Singer, Shahbuland Matiana, Joe Penna, and Omer Levy. Pick- a-pic: An open dataset of user preferences for text-to-image generation. Advances neural information processing systems, 36:36652-36666 2023.
Seul Lee, Karsten Kreis, Srimukh Prasad Veccham, Meng Liu, Danny Reidenbach, Yuxing Peng, Saee Gopal Paliwal, Weili Nie, and Arash Vahdat. Genmol: drug discovery generalist with discrete diffusion. In Forty-second International Conference on Machine Learning, 2025. URL https:: httpps/////peeeiee.....o net www.foooruuuuoreoooo
Xiner Li, Yulai Zhao, Chenyu Wang, Gabriele Scalia, Gökcen Eraslan, Surag Nair, Tommaso Bian- calani, Aviv Regev, Sergey Levine, and Masatoshi Uchara. Derivative-fre free guidance in continuous and discrete diffusion models with soft value-based decoding. CoRR, 2024.
Timothy Lillicrap, Jonathan Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint arXiv:
2015.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Representations, 2019.
Cheng Lu and Yang Song. Simplifying, stabilizing and scaling continuous-time consistency models. In The Thirteenth International Conference on Learning Representations, 2024.
Simian Luo, Yiqin Tan, Suraj Patil, Daniel Gu, Patrick von Platen, Apolinário Passos, Longbo Huang, Jian Li, and Hang Zhao. Lem-lora: A universal stable-diffusion acceleration module. arXiv preprint arXiv:2311.05556 2023.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, loannis Antonoglou, Daan Wier- stra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint arXiv:1312.560.. 2013.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Belle- mare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 518(7540):59-5333 2015.
11

---

Published as a conference paper at ICLR 2026
Ashvin Nair, Abhishek Gupta, Murtaza Dalal, and Sergey Levine. Awac: Accelerating online rein- forcement learning with offline datasets. arXiv preprint arXiv:2006.09:59, 2020.
Brendan O' D'Donoghue, Remi Munos, Koray Kavukcuoglu, and Volodymyr Mnih. Combining policy gradient and q-learning. In International Conference on Learning Representations, 2017. URL
/tttps://openreee
net forum? d= d=Bkidd61e1ekkkk
Dustin Podell, Zion English, Kyle Lacey, Andreas Blattmann, Tim Dockhor, Jonas Müller, Joe Penna, and Robin Rombach. Sdxl: Improving latent diffusion models for high-resolution image synthesis. arXiv preprint arXiv:2307.01992, 2023.
Mihir Prabhudesai, Anirudh Goyal, Deepak Pathak, and Katerina Fragkiadaki. Aligning text-to- image diffusion models with reward backpropagation. arXiv preprint avXiv:2310.03739, 2023.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pp. 8748-8763. PmLR, 2021.
Robin Rombach, Andreas Blattmann, Dominík Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEF/CVF confer- ence on computer vision and pattern recognition, pp. 10684-10695, 2022.
Tom Schaul, John Quan, Ioannis Antonoglou, and David Silver. Prioritized experience replay, arXiv preprint arXiv:1511.05995 2015.
Chrisoph Schuhmann. Laion aesthetics. URL https://taiin. aibllog/aioon-eesshitiiss 2022.
John Schulman, Xi Chen, and Pieter Abbeel. Equivalence between policy gradients and soft q- learning. arXiv preprint arXiv:1704.0644 2017a.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint 17X1v:1707.06347, 2017b.
Marcin Sendera, Minsu Kim, Sarthak Mittal, Pablo Lemos, Luca Scimeca, Jarrid Brotor-rrooks, Alexandre Adam, Yoshua Bengio, and Nikolay Malkin. Improved off-policy training of diffusion samplers, 2025. URL https: / arxiv. crt. abs
Joar Skalse, Nikolaus Howe, Dmitrii Krasheninnikov, and David Krueger. Defining and character- izing reward gaming. Advances in Neural Information Processing Systems, 35:9460-00 090009991,111 2022.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In Interna- tional Conference on Learning Representations, 202
Yang Song. Jascha Sohl-Dickstein, Diederik P Kingma, Abhishck Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In Interna- tional Conference on Learning Representations, 202
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. In International Conference on Machine Learning, pp. 3221 1-22552. PMLR, 2023.
Masatoshi Uehara, Yulai Zhao, Tommaso Biancalani, and Sergey Levine. Understanding reinforce- ment learning- based fine-tuning of diffusion models: A tutorial and review. CoRR, 2024a.
Masatoshi Uehara, Yulai Zhao, Kevin Black, Ehsan Hajiramezanali, Gabriele Scalia, Nathaniel Lee Diamant, Alex M Tseng, Tommaso Biancalani, and Sergey Levine. Fine-tuning of continuous- time diffusion models as retropy-r control. CoRR, 2024b.
Masatoshi Uehara, Yulai Zhao, Kevin Black, Ehsan Hajiramezanali, Gabriele Scalia, Nathaniel Lee Diamant, Alex Tseng, Sergey Levine, and Tommaso Biancalani. Feedback efficient online fine-tuning of diffusion models. arXiv preprint arXiv:2402. 16359, 2024c.
12

---

Published as a conference paper at ICLR 2026
Siddarth Venkatraman, Moksh Jain, Luca Scimeca, Minsu Kim, Marcin Sendera, Mohsin Hasan, Luke Rowe, Sarthak Mittal, Pablo Lemos, Emmanuel Bengio, et al. Amortizing intractable in- ference in diffusion models for vision, language, and control. Advances in neural information processing systems, 37:76080-761 2024.
Xiaoshi Wu, Yiming Hao, Keqiang Sun, Yixiong Chen, Feng Zhu, Rui Zhao, and Hongsheng Li. Human preference score v2: A solid benchmark for evaluating human preferences of text-to- image synthesis. CoRR, 2023.
Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. 2019.
arXiv preprint arXiv:
Jiazheng Xu, Xiao Liu, Yuchen Wu, Yuxuan Tong, Qinkai Li, Ming Ding, Jie Tang, and Yuxiao Dong. Imagereward: Learning and evaluating human preferences for text-to-image generation.
Advances in Neural Information Processing Systems, 36:15503-15935, 2023.
Ling Yang, Zhilong Zhang, Yang Song, Shenda Hong, Runsheng Xu, Yue Zhao, Wentao Zhang, Bin Cui, and Ming-Hsuan Yang. Diffusion models: comprehensive survey of methods and applications. ACM computing surveys, 56(4):1-39, 2023.
Richard Zhang, Phillip Isola, Alexci A Efros, Eli Shechtman, and Oliver Wang. The unreasonable effectiveness deep features as a perceptual metric. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 586-595, 2018
Xiangxin Zhou, Liang Wang, and Yichi Zhou. Stabilizing policy gradients for stochastic differential equations via consistency with perturbation process. In International Conference on Machine Learning, pp. 19-623333 PMLR, 2024.
13

---

Published as a conference paper at ICLR 2026
A KL-REGULARIZED RL IN DIFFUSION MDP
A.1 DERIVATION OF SOFT OPTIMAL POLICY
In this section, present the derivation of the soft optimal policy in KL-regularized RL for diffusion MDP in Section 3.2. Since we aim to maximize rewards while mitigating over-optimization, we define the objective follows (Uchara al., 2024a):
((-1, aD (po(: 2.)lle' mm))))
p' arg max E (T)
(17)
arg max E (7)
r(zo)
Σ DKL(Po(: m))ll' (
(18)
Pe
where (IT: .......... To) denotes the trajectory denoised by training policy pe and p' denotes reference policy (i.e. pre-trained diffusion model). So diffusion model py tries to maximize reward while staying close to the pre-trained diffusion model, which has naturalness and diversity. Then we define the soft Q-function Q (Te, x(-1) as the expectation of reward with the KL-divergence through the trajectory from the state-action pair It-1):
Q 2t-1) = R(F,, 1t-1) 1) + Bps
- (ool 2x) mz))) Iz, Ht-1
(19)
c-1 Dxs (oaa) 20, 21,
Exs r(zo)
Σ 1=1
(20)
By using the definition of soft Q-function, can reformulate our objective in Equation (18) as follows (Haarnoja et al., 2017):
T
p arg max
Euz,
[ (Q (It, Tt xt-1)) aD (p (1)))))))
(21)
Pe
We can define soft optimal Q function Q (Ite Vt 1) as expected return of taking action _;-1 from state and subsequently following the optimal policy p":
Qraa (Fe, x1) = 1xp r(zo) - a Σ (mml
tts xz-1
R=
(22)
By replacing the soft optimal function 22 with Equation (21), our objective simplifies to finding the policy po (-12;) that maximizes the expected soft optimal Q-function while remaining close to the reference policy p':
p" (-1xx)= = arg max E En, 1 (+1m,) (Qion (Zt, 21-1) - a oD (po (=)=)))))))))))))))
(23)
Po
Then by expanding the Equation (23) explicitly, we can re-express the objective maximization prob- lem KL divergence minimization problem (Haarnoja et al., 2018):
a log Po p x+)
p" (x)) = arg max Ez, 1-po((12) po())e))
Qinn(
(24)
// (x (it)
p p'axo-1 (F)) exp (Qien ( (Ie, xt-1)/a1 log
arg max Ez, Ex. (0p(1()())))) po()
(25)
( x()
arg DEL pol. xx) po(:)z,)
Z(x²) p'( xx)) exp ( Qioinltte
(26)
where Z(x.) = p' (x)) exp(Q2oft (Ie, (t-))
partition function which normal-
izes the distribution. Then the closed form soft optimal policy p² (-(ar)) is given by:
p'( xx) exp Qinm (Itz )/a)
p² (: xx)
(27)
J (2.1 xx) exp( ( (Ze, xt- xz= 1)/a) dx
14

---

Published as a conference paper at ICLR 2026
A.2 SOFT BELLMAN EQUATION
We begin by defining the soft optimal value function as follows:
VOn Von (Ize = Ep
Σ R(-1) (D a
(28)
= Bp
r(zo) - Σ (Hall
(29)
k=1
Then we can express the soft optimal Q-function in terms of the soft value function:
Qun (Es, [t-1) = R(I,, 2t-1) + Ep.
(R-1) aDxI (
ze. -1
R(xx, x1-1) + (
(30)
In the same way, we can rewrite the soft optimal value function as the soft optimal Q-function and the KL-divergence the current state Fz.
V (xt) = -us---ppr((
Epp
((() aDD (+ax))ll (1==)) Zes Ie-1
( [Qaa (Ie, x2-1)) aD (p' (1
(31)
By substituting the soft optimal policy to Equation (31), we can get the soft Bellman equation as follows:
Vo (x4e) = E. Be
(oin( - a log (3(
(32)
Eeg- (-p Qt (E (21, alo Z(r, (xm-1/a)) exp(Qin (en (x (2))
(33)
Esg 1mp (aai) a log (Ixa) log ((xe)
(34)
log / (x mz) exp((iio((z x1-1)/a) dxr-1
(35)
R(Ie, D1 11 + Vina Vinn (26-11)
a log ( expl
(36)
A.3 APPROXIMATE, SOFT Q-FUNCTION USING TWERDIE'S S FORMULA
Training value network is often challenging in diffusion MDP (Hu al., 2025). Using posterior mean estimation, fo(xx) obtained by Tweedic's 's formula (Chung et al., 2023): black
(o(x)) Ey [xolxx]]
Vou (x+ + v1 ãe e (1r, f))
(37)
can approximate the soft optimal Q function (Li al., 2024):
Dion = V (x
(38)
= a. log Eug.... 2 o (-10-1)) exp R(11,
84-1
(39)
22 =r(20(2:-1).
(40)
Then we have an approximation of the soft optimal Q function in Equation (27). Therefore, by minimizing the KL-divergence in Equation (26), corresponding loss function is (Haamojj et al., 2018):
get surrogate soft optimal policy. The
L(0) Exr. [m po (-1x))- Q r.-1) + (Dr (xx- + a au (po(0))))) (=
(41)
es Ex. Then the gradient of the Equation (42) is given by: (oC(O) = Ex, [B.N(o, a._ ( Vouu(Zi, t) (D
(42)
(43)
with the reparameterization #t-1 n0( 1) ote
15

---

Published as a conference paper at ICLR 2026
B KL-REGULARIZED RL IN DISCOUNTED MDP
Considering the inherent stochasticity of the diffusion MDP, introduce discount factor 7 [0, 1) to improve credit assignment. Then we can rewrite the original KL-regularized RL objective Equation (18) follows:
Σ 7 T-1 (((1), - aDxa xx)llp' (: mr))
p- = arg max E Euyppo
(44)
Pe
T
arg max E. (pp(tr
- -T-))()))
- ----' aDKilpp I ax)ll'( xx)))
(45)
Pe
1=1
The discount factor causes the signal from the terminal reward r(3o) decay exponentially as increases, downweighting the contribution of early-stage actions. Following the definition Equa- tion (22), we define the soft optimal Q-function for (Iz, 2t-1), considering the discounted return:
(in( = r( x1- + Rpp
Σ 1
- - aD (aall (-xxx)(())x)))
34: Tt-1
(46)
[-1
R(2,, Eyp (((0))))
x (H (((x))lll(((+))))))) Jes xe-1
(47)
The derivation in (21 23) simplifies the global objective into per-step maximization problem by using the soft optimal Q-function. This derivation still holds for our discounted MDP, yielding the same soft optimal policy
p(( x2) exp (Giot (xt, )/a)
p' 1 xz)
(48)
I pl (x 24) exp( (4 dxe-1
B.1 SOFT BELLMAN EQUATION IN DISCOUNTED MDP
Again, we define soft value function with discounted return follows:
Vi (x+) = Ep- Σ T" t
(R(xx, - a Dxi(p"(: xR) 1| (-1 xx))
xt.
(49)
= Epp
tha Dxu(p"(11 xx) (. xx))
(50)
k=1
Then we can derive the soft optimal Q-function in terms of the soft optimal value function:
Ounlx1, = H(xx, 26 11) + Ep-
((( aDxu (-low)) xe 21
(51)
= R(xx, 2x-1) + Ep
Σ (t- (MEx, aDxL (p"((pxx))l (-1zx))
(52)
-(2x, + Vian (F
(53)
16

---

Published as a conference paper at ICLR 2026
Following the derivation (3), get the soft Bellman equation follows:
(2s) Brrup" (0-7((()) x DRalp" (- aa)) 1| mw)) 22
(54)
k=l
=Bu (opp-rpp(a Brup"
? ((x0) Exe DRIlF" (: xx)) 11 1' (0(11x)))) xxx)
24, 31-1
km
(55)
-Bu (1 (1 (xa D s4) p' zto)
(56)
M(as, xx-1) + y Vein (1)-1)
aa Br, (o))))
exp
(57)
B.2 BOUNDS OF SOFT Q-FUNCTION
In this section, we present the detailed derivation for the bounds of the soft optimal Q function:
a alog E Pxo exp r'= r(wo)) S Qiot (F x1-1 s a log E a (
[(zo]
exp
We can rewrite the soft Bellman equation in Equation (57) as:
V cxp Wonn(a (3T)
V son (2z-1))
exp
(58)
The discount factor E 1) makes the function u + u² concave. Consequently, each recursive expansion of the soft Bellman equation invokes Jensen's inequality, creating gap that yields bound on the soft optimal Q-function, unlike the direct expansion in Section A.3. To derive the
bounds, we first denote Z, (z,)) = (x)nVVa Then we can rewrite the Equation (58), and it's one step expansion as follows:
Z.(xe) (n Z (
(59)
Z Ze-1 (20-(1 Ee 20p 1) (1-22-2220 2)
(60)
By Jensen's 's inequality:
ZI-1(x=-1)" (Bn. (- (m, (
E En (< Z02((002))
(61)
By applying Equation (61) to Equation (59):
Ze (xx) Ez 10-- 7 (x1 (x]-1)
(62)
( Ir (11
Zz 20=2(00-))
(63)
- Zz 202
(64)
Then we can recursively expand Jenson's 's inequality and get the lower bound of Zz(xx):
(x( 2 2m Z 2
(65)
(66)
E Ex,, 2
Z (1 $1
(67)
2 ((0))) a
-E (
exp
(68)
Since Qiaa (xe, 21-1) Von( can derive the lower bound of the soft optimal Q function:
Qiaa (Me; xz-1) = Viin (ax-1)
(69)
- - (
(70)
exp (2-
2 T Eo..... 201-////
(71)
a
17

---

Published as a conference paper at ICLR 2026
Now, also derive the upper bound of the soft optimal of the function using the same procedure as 59 71. By recursively expanding Jenson's inequality in the opposite way, we can obtain the upper
bound of (zx) as follows:
Ze(xe) (m- Zi (x
(72)
( (1 )
(73)
(74)
B (11oe) 21 (2x1))
(75)
exp (exp(((aa)
(+bo))
(76)
Then the upper bound of the soft optimal Q function is defined as:
Qion (It, 14-1) = > [
(77)
=
(78)
(1(Fo)
< : E... 2 2---0pp-1aaa- (-11--1)
exp
(79)
We have thus derived lower and upper bounds on the soft optimal Q-function:
log Soonsrrp'l [exp( 2-- r(xa))] & Qian (Fe, #0-1) S - log E 2 2mp' (1be
rizo] exp
a
(80)
B.3 APPROXIMATION OF SOFT OPTIMAL FUNCTION IN DISCOUNTED MDP
We observe that applying Tweedie's formula to both the lower and upper bounds of the soft Q- function yields the identical first-order approximation:
y a log Bra... 20p-2 exp 2 )r(oo) )| ] ~ xx- "((r(((((r-i)))
(81)
exp (1(20) 7-- ((00(((((())1)))
7 7-1 a log Egg -- (+
(82)
Therefore, we can approximate our Qun (We, I2-1) as follows:
V (Fe, x1-1) xo 1
(83)
18

---

Published as a conference paper at ICLR 2026
C VISUALIZATION OF REWARD OVER-OPTIMIZATION PHENOMENA
Aesthetic:
HPS 0.277
ImageReward 0. 745
- 0.648
DivamSii 0.755
Aesthetic 0274
InagdReward 0.847
LPIPS 565
DreamSim: 0.608
Aestbetic
IewgeReward -1.303
L 0.416
DivamSin: 0.378
Aestaticc 8.2
IPS:0223 0.37
ImageRoward
LPIPS: 0.452
DreamSim 0.325
Figure 8: Visualization of reward over-optimization during fine-tuning. Each row displays gener- ated images from four sequential training checkpoints, arranged chronologically from top to bottom. As aesthetic scores increase throughout training, the images exhibit two critical failure modes: mantic collapse (loss of alignment with prompts) and diversity collapse (convergence toward similar abstract patterns). Quantitative metrics evaluated each checkpoint are shown the left.
In Figure we illustrate two characteristic phenomena of reward over-optimization: (1) semantic collapse and (2) diversity collapse. The examples generated from model fine-tuned with the DRaFT-1 method at four different training checkpoints, while the corresponding training curves are shown in Figure 3.
Semantic collapse. As training progresses (top to bottom), images with higher aesthetic scores progressively lose alignment with their original prompts, dissolving into abstract textures. This phenomenon is reflected in the decline of alignment metries such HPS and ImageReward.
Diversity collapse. At the same time, images with higher aesthetic scores lose diversity, with backgrounds and textures converging toward highly similar patterns. This phenomenon is reflected the decline of diversity metrics such as LPIPS and DreamSim.
19

---

Published as a conference paper at ICLR 2026
D EXPERIMENT DETAILS
D.1 IMPLEMENTATION DETAILS
For all experiments, we use Stable Diffusion v1.5 (Rombach et al., 2022) as the base model and adopt LoRA (Hu et al., 2022) for compute-efficient fine-tuning. For the consistency model, employ LCM-LoRA (Luo et al., 2023) on Stable Diffusion v1.5. We set the diffusion sampling steps to 50. All experiments image fine-tuning and on two NVIDIA A6000 48 GB GPUs for online black-box optimization. For SQDF, fine-tuning takes approximately 62 and 401 per update step when using aesthetic and HPS as reward functions, respectively, and 490 per inner update step for online black-box optimization.
conducted two NVIDIA RTX 4090 24 GB GPUs for text-to-
D.1.1 FINETUNING FOR TEXT-TO-IMAGE DIFFUSION MODELS
For all baselines on text-to-image tasks, we use the A optimizer (Loshchilov & Hutter, 2019) parameters with 81 0.9, 32 0.999, weight decay of 0.0001, and set the classifier- free guid- ance (CFG) scale to 5.0 for all methods. Following (Clark et al., 2023), we divide our experiments into two distinct settings. For our large-scale experiments, which involve fine-tuning on the HPS dataset, we used a batch size of 258, LoRA rank of 32, and 500 training steps. For the small-scale experiments, the aesthetic score, we used a batch size of 64, a LoRA rank of 4, and 2000 training steps, except for DDPO, where we use 500 training steps, which takes the same wall time SQDF. We provide implementation details for all experimental methods below. To ensure fair compari- and reproducible results, used official codebases and auchor-eecommended hyperparameters whenever available.
SQDF For our implementation of SQDF, we utilized DDPM sampling with 50 steps. In small- scale experiments, we set the discount factor and the learning rate to le - 3. In large-scale experiments, these values were adjusted to = 0.93 and a 0.05, and the learning rate was set to 5e 4, respectively.
0.9, the KL-regularization coefficient to 2,
DRaFT
As there is no official implementation available for DRaFT, we use the AlignProp code- base, methodologically similar concurrent work. We utilized DDIM sampling with steps. The original paper notes that there is no significant performance difference between DDIM and DDPM samplers. Following the settings outined their paper, we used a learning rate 4e for small-scale experiments and 2e- 4 for large-scale experiments.
ReFL We implemented ReFL on top of the AlignProp codebase. We used DDPM sampling with 50 steps. The original ReFL method generates samples by denoising to a random timestep t and then performing one-step prediction. To adapt this to our 50-step diffusion process, we set by sampling it uniformly from the range 50). For the small-scale setting, we used a learning rate of le 5. As the original paper does not provide hyperparameters for large-scale setting, and we observed slow convergence with the smaller learning rate, we used 2e to match the DRaFT setup.
DDPO For DDPO, we utilized the official codebase provided by the authors. We employed DDPM sampling with 50 steps. We used learning rate of 3e for both small-scale and large-scale exper- iments. This value was chosen it is the recommended setting for reproducing results in the official GitHub repository. Since no specific hyperparameter was provided for large-scale experiments, we extended this recommended setting to that context well.
KL-regularized baselines For DRaFT+KL and DDPO+KL, we introduce KL-regularization into the original objective same as our method:
T
p" arg max E Pe
r(av) - a Σ Dxi(pol- 12) 22))
Meanwhile, the learning method differed between the two methods. For DRaFT+KL, we calculated the KI. divergence each denoising step and added this value the reward signal. We set the
20

---

Published as a conference paper at ICLR 2026
default regularization coefficient to 0.04, value chosen to match the expected value of the KL term in SQDF. All other hyperparameters for DRaFT+KL were kept the same as the original DRaFT.
For DDPO+KL, we implement KL regularization by computing the KL divergence at each denoising timestep, detaching it from the computational graph, and subtracting it from the advantage before policy updates. We use a regularization coefficient of a = 0.2, with all other parameters following the original DDPO configuration.
D.1.2 ONLINE BLACK-BOX OPTIMIZATION WITH DIFFUSION MODELS
For fine-tuning the diffusion model, we used a learning rate of le with the AdamW opti- mizer (Loshchilov & Hutter, 2019) of 81 0.9, A = 0.999, and a weight decay of 0.1, and set the classifier-free guidance scale to 7.5. For training the reward proxy, we used MIP top of a frozen CLIP encoder (Radford et al., 2021). To incorporate an uncertainty bonus, we imple- mented techniques, (1) Bootstrap: statistical ensemble of neural networks, created by resam- pling the given dataset, and (2) UCB bonus derived from the final neural network layer. We adopt the number of bootstrap head and UCB parameter C1 0.01, 0.001.
SQDF For our implementation of SQDF, we set discount factor 7 0.9 and KL- regularization coefficient to a = 1. For each outer loop, we set inner iteration=20 and samples per iteration 64. We adopt prioritized replay buffer for training, detailed in Section E.
SEIKO We adopted the hyperparameter settings from the official codebase for SEIKO (Uchara al., 2024c). Specifically, we set the Kl-regularivation coefficient to a 0.01 and performed inner iterations with 64 samples per each. The truncated pack-propagation step K was uniformly sampled from (0,50), where K limits the number of steps through which gradients are backpropa- gated.
PPO+KL We implement KI.-penalized PPO (Schulman et al., 2017a) as mentioned in Sec- tion D.1.1 and train directly using the oracle reward with the same budget. We employ the hyper- parameter settings from the official DDPO codebase, adding the KL-regularization with a 0.2.
D.2 TASK AND PLOTTING DETAILS
In this section, we provide a detailed explanation the tasks, along with corresponding figures that offer further clarification for the interpretation.
D.2.1
TOY TASK
The toy example in the second row of the Figure 2 demonstrates intuition of employing the consistency model. We use a 2D Gaussian Mixture Model (GMM) as the ground truth data distribu- tion. GMM contains nine components with means positioned
regular 3 3 grid spanning from (-4, 44) to (4, 4), where each component has an isotropic covariance of 0.3. We visualize one-step predictions clean sample (t from various noise levels (t. 35, 25, 15).
D.2.2 TEXT-TO-IMAGE FINE-TUNING
For both tasks in Figure 3, we plot six points for each method, using checkpoints taken at uniform intervals up to defined terminal rounds. While the aesthetic score is independent of the alignment and diversity score, HPS has a positive dependency on the alignment score. To this end, experi- ment and plot Figure with different criteria. For the top row (optimizing for aesthetic score), the terminal epoch was set as the point at which each method first reached target reward of 8.0, where baselines exhibit severe semantic collapse. If not achieved, we report the maximum reward. For the bottom row (optimizing for HPS), the terminal epoch was set the maximum training epochs of 500 for all methods.
KL-regularized baselines comparison In Figure 4, each data point on the trade-off curves rep- resents the final performance of model finetuned with specific KL-regullrization coefficient a..
21

---

Published as a conference paper at ICLR 2026
Darker points correspond to larger values of For SQDF, started with the default 2 and increased to 3 and 4 to depict performance at lower target reward points. For DRaFT+KL, we started with 0.03, then increased the coefficient to 0.035 and 0.04. For DDPO+KL, ducted an empirical search and plotted three points using values of [0.2, 0.3, 0.4). We note that, different baselines have different scales of parameter settings that work well, we adjust best as possible for the other baselines.
D.2.3 ONLINE BLACK- BOX OPTIMIZATION
Following (Uehara et al., 2024c), our online black-box optimization task employs an iterative pro- cedure, which is repeated over outer loops. We assume the Aesthetic score a black-box oracle reward function. Each loop consists of steps: (1) sampling an image from the current policy, (2) querying the oracle reward function, (3) updating the reward proxy model with oracle feedback, (4) fine-tuning the policy with the updated proxy model. We execute this process with [1024, 2048, 4096, 8192] oracle queries per loop, for total of 15360 feedback. After the final loop, we port the optimization target reward (Aesthetic Score), unseen rewards (ImageReward, HPS), and diversity metries in Table 1.
D.3 TRAINING AND EVALUATION PROMPTS DETAILS
For optimizing the aesthetic score, use the 45 simple animals dataset training and evaluation prompts, following Black et al. (2023). For optimizing HPS, we use the HPDv2 dataset (Wu et al., 2023), which contains 3,200 prompts divided into four categories (Animation, Concept Art, Paint- ing, and Photo), with 800 prompts per category. We split the prompts into training and evaluation prompts using fixed random seed, sampling 12 evaluation prompts per category and using the maining prompts for training. For online black-box optimization, we train on the 45 simple animals dataset and evaluate on another six prompts (snail, hippopotamus, cheetah, crocodile, lobster, and octopus), following the approach of (Uchara et al., 2024c).
D.4 EVALUATION DETAILS
For evaluation, we generate 32 images per prompt from the evaluation set. To measure alignment, all sampled images are scored using the ImageReward, HPSv2, and pickscore, and report the mean score for each metric. To measure diversity, we compute pairwise distances across all generated images. Specifically, for LPIPS, use the perceptual distances directly produced by the model and report the mean of all pairwise values. For DreamSim, we first extract feature representations for all images and then compute the average pairwise cosine distance between these features.
22

---

Published as a conference paper at ICLR 2026
E ADDITIONAL ANALYSIS
E.I OPTIMIZE WITH PRIORITIZED REPLAY BUFFER
One of the key innovations of SQDF is off-policy update with a replay buffer. In this subsection, analyze the effect of using prioritized replay buffer (Mnih et al., 2015; Schaul et al., 2015) on overall performance.
We prioritize samples that from high-reward trajectories and closer to the clean sample. Specifically, for each sample Te in the buffer, we assigned value of y'r, where r is the reward obtained from the trajectory's 's final state (20) and / is the timestep. Training samples are then drawn from the buffer according to probabilities derived from these priority values.
(a)Eiocch vs Acsthetic
(b) Aesthetic vs
Aesthetic LPIPS
(d) Aesthetic vs DreumSim
0.44
Aesthetic
2.-
LPIPS-DR DW
DroamSinn-Di
6.5
0:5
0.4
335
0271
6.2
26
169
65
24
epoch
Aesthetic
Aesthctic
Aesthetic
SQDF w/o Prioritization
SQDF Prioritization
Figure 9: Effect of Prioritized buffer Aesthetic Score.
Figure shows the training dynamics of SQDF with and without prioritized buffer when optimiz- ing for the aesthetic score. For plots (b)-(d), we use checkpoints taken at uniform intervals up to the epoch where each method reached an aesthetic score of 8.0.
The results show that SQDF with Prioritization converges to high aesthetic scores more rapidly, but exhibits trade-off by achieving lower alignment and diversity scores in the high-reward range (aesthetic score 7.0). This analysis indicates that prioritized buffer improves sample efficiency and suggests that integrating more advanced prioritization techniques promising direction for further performance enhancements.
E.2
EFFECT OF VARYING DISCOUNT FACTOR
(a)Epoch vs Aesthetic
(b) Acsthetic HPS
(c) Aesthetic LPIPS
(d) Aesthetic DreamSim
0,29
Aeshheic
2.5
LPPPS- Div
0:40
DreomSimm
6
0:55
69
6d
53
spuch
Aesthetic
Aesthetic
Aesthetic
40.90
0085
0.95
Figure 10: Training curves with different 7. A higher discount factor leads to faster convergence but with less diversity and alignment score.
We investigate the effect of varying the discount factor y and provide recommendation for its selection. We conduct experiments with y E {0.85, 0.9, 0.95] under a diffusion denoising step setting of T 50, using the aesthetic score the target reward.
As shown in Figure 10, larger values of y lead to faster convergence and higher optimization of the target reward, but at the cost of reduced alignment and diversity scores. Conversely, smaller values
23

---

Published as a conference paper at ICLR 2026
of y result in slower convergence, but the less aggressive reward optimization allows the model to achieve better diversity and alignment. These results suggest that the value of y controls a clear trade-off between optimization speed and sample quality.
24

---

Published as a conference paper at ICLR 2026
E.3 QUALITATIVE COMPARISON ON ONLINE BLACK- BOX OPTIMIZATION
Pretrained Reward: 5.27
SQDF-UCB Reward: 8.09
SQDF-Bootstrap Reward: 8.11
SEIKO-UCB Rewand: 7.65
SEIKO- SEIKO-Bootstrap Reward: : 7.77
PPO+KL Rewand : 6.32
Figure 11: Qualitative comparison of online black-box optimization. Samples generated by models fine-tuned with different methods using the same oracle query budgets. Prompts: ["Cheetah", "Oc- topus", "Hippopotamus", "snail", "Crocodile", "Lobster"].
25

---

Published as a conference paper at ICLR 2026
E.4 QUALITATIVE COMPARISON ON HPSv2
Qualitative comparisons are presented in Figure 12, where HPS used as the target reward. Al- though SQDF was optimized for slightly lower target reward compared other methods, it sistently produced samples that are well aligned with the input prompts. For comparison, ReFL. achieved higher target reward but showed relatively less prompt consistency. These results indi- cate that SQDF can effectively fine-tune under the HPS reward.
Pretrained
Ours
DRaFT
ReFL
DDPO
Prompt: portrait of spacesuit, with backdrop. Krenz Cusbart, popular and design platforms.
Prompt: The imuge stylized medium shot portruit of business monkey that would be suitable for box advertising.
sitting g bench with behind her. Prompt:
Prompt: : man riding on the back of a motoreycle down a highway,
Figure 12: Qualitative comparison of different finetuning methods on the HPS.
26

---

Published as a conference paper at ICLR 2026
E.5 QUALITATIVE COMPARISON WITHIN PROMPT SAMPLES
In this subsection, we visually evaluate the diversity of generated samples within individual prompts. As shown in the Figure 13 17, SQDF produces diverse images for the same prompt, demonstrating strong within-prompt variation. In contrast, other methods tend to generate visually similar outputs, indicating lower diversity. These results confirm that SQDF effectively preserves sample diversity and naturalness while optimizing for high target rewards.
Pretrained
Rowrd-527
Ours Roward 7.77
ReFL Reward 2.65
DRaFT Rowsrd 803
KKL Reward 6.12
DDPO Reward 6,46
DDPO+KL Reward 6.75
Figure 13: Samples generated by models fine-tuned with different methods for the prompt: "Bee".
27

---

Published as a conference paper at ICLR 2026
Pretrained Reward 5.05
Ours Reward 8.16
ReFL Reward 8.17
DRaFT Reward 7.97
DRLFT+KL Rowad 6.99
DDPO Rewar 2.24
DDPO+KL Rewar: 6.54
Figure 14: Samples generated by models fine-tuned with different methods for the prompt: "Cat".
28

---

Published as a conference paper at ICLR 2026
Pretrained Reward3311 5.19
Ours Rewar 7.90
ReFL Reward 7.56
DRaFT Reward 8.15
KRL Rowad: 6.39
DDPO Rewar 2001
DDPO+KL Rewar: 6.34
Figure 15: Samples generated by models fine-tuned with different methods for the prompt: "Turtle".
29

---

Published as a conference paper at ICLR 2026
Pretrained Reward 0.26
Ours Roved 0.30
DRaFT Rewar 0.32
ReFL Roved 0.33
DDPO Reward 0.88
Figure 16: Samples generated by models fine-tuned with different methods for the prompt: "A landscape with an art nouveau building".
Pretrained
Ours 0.31
DRaFT Rewar 0.31
ReFL Rewad 031
DDPO Rewad- 0.27
Figure 17: Samples generated by models fine-tuned with different methods for the prompt: "The image is a stylized medium shot portrait business monkey that would be suitable for box art or advertising".
30

---

Published as a conference paper at ICLR 2026
F EXTENSION TO ADVANCED BACKBONE
In this section, we further extend our experiments by replacing the base model from SD v1.5 to SDXL (Podell al., 2023) to evaluate the effectiveness of SQDF a more advanced and larger- scale architecture. All experiments are conducted using three random seeds.
First, we fine-tune SDXL using both DRaFT-1 and SQDF optimize the aesthetic score. As shown in Table 3, SQDF outperforms DRaFT, achieving superior aesthetic scores and alignment while successfully maintaining diversity.
Metric
SDXL (Pre-trained)
DRaFT
SQDF
Aesthetic (1)
5.45
7.18
7.86
IeageReward (1)
0.88
0.91
1.21
HPS (t)
0.28
0.27
0.28
LPIPS-Div (1)
0.59
0.49
0,51
DreamSim-Div (1)
0.72
0.48
0.57
Table 3: Comparison of DRaFT and SQDF on fine-tuning SDXI for aesthetic score.
Furthermore, to investigate whether the effectiveness of SQDF depends on the size of the base model, we compare the relative performance improvements between SD 1.5 (860M) and SDXL (2.6B). As shown in Table the relative improvement achieved by SQDF is highly consistent across both architectures.
Metric
SD1.5
SD1.5 SQDF
%
SDXL
SDXL + SQDF
%
Aesthetic (f)
5,46
7.87
44.14%
5.45
7.86
44.22%
IeageReward (f)
0.93
1.14
22.58%
0.88
1.21
37.50%
HPS (T)
0.28
0.28
0.00%
0.28
0.28
0.00%
LPIPS-Di Div (1)
0.65
0.56
-13.85%
0.59
0.51
-13.56%
DreamSim-Div (f)
0.75
0.58
-22.67%
0.72
0.57
-20.83%
Table 4: Relative performance improvements of SQDF on SD1.5 and SDXL for aesthetic score.
These results demonstrate that SQDF optimizes the target reward while mitigating over- optimization, regardless the underlying diffusion backbone.
31

---

Published as a conference paper at ICLR 2026
G UNCURATED SAMPLES
We present uncurated samples generated by models fine-tuned with each method discussed in the main paper.
Pretrained 3.3
Ours Reward 804
ReFL 730
DRaFT 8.06
KRL Rowad 2.17
DDPO Rowed 2.18
DDPO+KL 2.28
Figure 18: Uncurated samples generated by models fine-tuned with different methods for the prompt: "Frog".
32

---

Published as a conference paper at ICLR 2026
Pretrained Reward 5.37
Ours Reward 7.82
ReFL Reward 7.60
DRaFT Reward 8.04
KKL Reward 6.02
DDPO Reward 6.48
DDPO+KL Rewar: 6.34
Figure 19: Uncurated samples generated by models fine-tuned with different methods for the prompt: "Bee".
33

---

Published as a conference paper at ICLR 2026
Pretrained Reward 0.28
Ours Rewad 0.34
DRaFT Reward 0.34
ReFL Roward 0.35
DDPO Reward 0.29
Figure 20: Uncurated samples generated by models fine-tuned with different methods for the prompt: A portrait of cat in spacesuit, with surreal backdrop, by Krenz Cushart, popular on art and design platforms.".
Pretrained Rowad 0.27
Ours Reward 0.31
DRaFT Reward 0.32
ReFL Roved 0.32
DDPO Reward 0.27
Figure 21: Uncurated samples generated by models fine-tuned with different methods for the prompt: "The image stylized medium shot portrait business monkey that would be suit- able for box advertising".
34

---

Published as a conference paper at ICLR 2026
H PSEUDO CODE OF SQDF
In this section, we provide pseudo codes of SQDF of the text-to-image fine-tuning task Algo- rithm and online black box optimization task Algorithm 2.
Algorithm Soft Q Diffusion Finetuning (SQDF)
1: Input: Pretrained diffusion model p': Trainable diffusion model pe; Replay buffer D; Training epochs N; Batch size B; Reward model re: Consitency model fø:
2: Initialize po deepcopy(p)) replay buffer D <- 0
3: for == :......,N NN do
4:
On-policy Sampling 5: for = 1,..,..,, B do
6:
Sample noise 2y (0,
7:
Run diffusion denoising steps using pe to obtain trajectory {ar, -T-1.... ....... 21
8:
For each (Iz, 1), store into D
9:
end for
10:
Policy Update
11:
Sample a batch ()) from D.
12:
for j- 1..... B do
13:
Perform one-step denoising o/-1 (=(2)-1lz(,,t)
Estimate the clean sample aif via consistency model f (x)_1)
15:
Compute reward r rr (
16:
Compute the loss & (Eq. 10)
end for
18:
Update 0 by minimizing I. B E, e
19:
end for
35

---

Published as a conference paper at ICLR 2026
Algorithm SQDF with Online Black-Bo Box Optimization
1: Input: Pretrained diffusion model p': trainable diffusion model ps; replay buffer D; Con-
sitency model Jøi outer rounds N; inner updates per round M; oracle query schedule (0,,02,-- ......-.. ON batch size B; surrogate reward model fą: oracle reward r* 2: Initialize < deepcopy(p") replay buffer 'D <- 0, labeled dataset (
3: Collect O2 initial samples {zbti-i using query oracle rewards {v/ 1 and store data in S
4: Train surrogate fgg on S {2) v²)
5: for k - N do
6: for m .., M do
7:
On-policy Sampling 8: for b = 1, ..... B do
9:
Sample noise IT ~ 10, Run diffusion denoising under pa obtain trajectory {ær, JT -----11--- . Store each state tuple (z₁, t) in D
11:
end for
13:
Policy Update
14:
Sample a batch [xx] 10) from D for l, do
15:
Perform one-step denoising ri-11 ~ po xx,11) 1)
17:
Estimate clean sample dig via consistency model fa(((_1)
Compute surrogate reward rJ (*t)
Compute the loss Ø (Eq. 10)
20:
end for
Update by minimizing L Eg l
end for
23:
Oracle Querying and Surrogate Update
for 1,.... O2 do
25:
Sample noise
(O,
26:
Run diffusion denoising under Po to obtain clean sample æo
27:
Query the oracle to obtain reward r* (zo)
28:
Store (zo. 3) in
29:
end for
30:
Retrain surrogate is on the dataset S
31:
end for
36


Reviews:
[{"review_id": "Ee4U8oSxRN", "reviewer": "Reviewer_M6wy", "summary": "- Problem: when you RL diffusion models it overoptimizes the reward and mode collapses and stops producing coherent images.\n- Solution:\n  - KL regularization to the original model (also common in past work)\n  - Soft-Q function (this is a normal Q-fn with an entropy term)\n  - Discount factor for better credit assignment\n  - Consistency models for more stable RL training\n  - Off-policy replay buffer - prevent mode collapse\nAlgorithm works like this:\n  - Generate samples from diffusion model, store them in the replay buffle\n  - Sample one of the partially noised samples, denoise it 1 step.\n  - Then take the consistency model and denoise the rest of the way to a clean image.\n  - Run the clean image through the RM.\n  - Backprop using policy gradient.\nExperiments\n  - 3 things they want to show (a) still good samples (b) less semantic collapse (c) less diversity collapse.\n  - To measure semantic, they use pretrained prompt-image alignment models like ImageReward.\n  - To measure diversity, they look at the mean pairwise distance as computed with 2 measures: LPIPS and Dreamsim features cosine similarity.\n  - Their experimental results look pretty legit:\n     - Show there’s an alignment-diversity tradeoff, but their method expands the pareto frontier\n     - Shows that there is an aesthetic-alignment tradeoff, but they minimize and ~almost remove it. They are slightly better than ReFL and much better than the other methods.\n     - Shows there is an aesthetic-diversity tradeoff and they’re way better than other methods.\n  - Ablations show the discount helps a bunch. It looks like the buffer helps with diversity (but slightly hurts reward) and the CM helps with reward (but slightly hurts diversity). But maybe these differences are small enough it’s noise?", "questions": "- In Fig 3, how did you get the diff points on each curve?\n- Table 1 - what is the number in parens. (confidence interval? standard deviation?)\n- Are the qualitative results randomly sampled or cherry-picked?\n- I suggest moving Algorithm 1 to the main paper, it was very helpful in understanding the algorithm.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "- Substantial baselines - they compared against several past works. They tested whether a simple addition to baselines (adding KL regularization) would improve them, and found that their method was still superior.\n- Ablations - removed each component they added to confirm that it contributes to the final result.\n- Clarity - the paper is clear and easy to read. The diagrams make sense.\nSignificant improvement above prior work, which is visible both in the qualitative results and the graphs.", "weaknesses": "See “questions” section for more details.\n\nIt is not clear to me that the minor improvements in diversity produced by the buffer (.01 to .02 in Table 2) is worth the extra complexity adding it to the algorithm requires and the hit to aesthestic score.\n\nSimilarly it seems like without the consistency model diversity improves (though this is a minor effect, the improvements to aesthetic score I think suggest the CM is still a good contribution to the algorithm).", "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": "9YAAUNGZ3v", "reviewer": "Reviewer_TRJT", "summary": "This paper focuses on fine-tuning diffusion models with downstream objectives while mitigating reward over-optimization. The proposed method solves KL-regularized off-policy RL with stabilized value estimation for controlling reward-diversity tradeoff. The method is compared with prior fine-tuning works on image generation with various rewards.", "questions": "1. There is no 'Q-learning' in the method since it approximates the value function using a consistency model. The name of the method (Soft Q-learning for Diffusion Finetuning) may be misleading.\n2. Using a frozen consistency model can introduce a different type of approximation bias due to the mismatch between the posterior mean of the fine-tuned diffusion and the consistency model. Why not fine-tune the consistency model together?\n3. What's the purpose of the online black-box optimization experiment? The paper lacks an explanation of which aspect of SQDF this experiment is trying to demonstrate.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Reward over-optimization has been a well-known problem in diffusion fine-tuning, and the paper proposes a well-structured method with results to support their claims.\n2. Each component of the method is well-motivated, and evidences are provided through ablation studies.\n3. Both quantitative and qualitative results show better reward-alignment / reward-diversity compared to prior works.", "weaknesses": "1. Lack of fine-tuning baseline: [1] have already proposed RL fine-tuning that focuses on mitigating over-optimization.\n2. Lack of training-free baselines: [2] has shown that training-free methods can mitigate over-optimization compared to fine-tuning, and [3], [4] evidence soft value function can also be used for training-free methods.\nWithout comparison with these methods or justification statements, it's unclear why fine-tuning is necessary.\n\n[1] Zhang, Ziyi, et al. \"Confronting reward overoptimization for diffusion models: A perspective of inductive and primacy biases.\" arXiv preprint arXiv:2402.08552 (2024).\n[2] Kim, Sunwoo, Minkyu Kim, and Dongmin Park. \"Test-time Alignment of Diffusion Models without Reward Over-optimization.\" The Thirteenth International Conference on Learning Representations.\n[3] Li, Xiner, et al. \"Derivative-free guidance in continuous and discrete diffusion models with soft value-based decoding.\" arXiv preprint arXiv:2408.08252 (2024).\n[4] Uehara, Masatoshi, et al. \"Inference-time alignment in diffusion models with reward-guided generation: Tutorial and review.\" arXiv preprint arXiv:2501.09685 (2025).", "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": "J4L0Cys7yv", "reviewer": "Reviewer_3B3B", "summary": "The paper introduces Soft-Q Diffusion Finetuning (SQDF), a novel method for aligning diffusion models with downstream reward functions while mitigating over-optimization. The core innovation is a training-free approximation of the soft Q-function using a single-step posterior mean estimate, which enables direct gradient-based optimization without the instability of value network training or high-variance Monte Carlo estimators. Experiments on text-to-image tasks demonstrate that SQDF achieves superior reward-diversity trade-offs compared to existing RL-based and direct backpropagation methods.", "questions": "Why did you choose the two-model architecture instead of directly fine-tuning the consistency model? Given that consistency models can generate samples using multiple steps, did you explore whether similar reward-diversity trade-offs could be achieved more efficiently by fine-tuning the consistency model directly?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The training-free soft Q-function approximation eliminates the need for unstable value network training.\n\n2. The paper demonstrates effectiveness across multiple tasks (aesthetic scoring, human preference optimization, black-box settings) with thorough comparisons against relevant baselines and ablation studies validating each component.\n\n3. Results show SQDF achieves better Pareto frontiers, optimizing target rewards while maintaining significantly better alignment and diversity metrics compared to methods that suffer from semantic and diversity collapse.", "weaknesses": "1. While using the consistency model to estimate the future value is intriguing, it also makes one wonder what about directly fine-tuning the consistency model, which could potentially achieve similar results more efficiently given its single-step generation capability. Also see Q1. \n\n2. While methodologically sound, the paper provides little discussion of computational overhead, training time comparisons, or memory requirements.", "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": "Ysa0xsamFH", "reviewer": "Reviewer_5VJm", "summary": "This paper introduces Soft Q-learning for Diffusion Fine-tuning (SQDF), a reinforcement learning–based framework designed to fine-tune diffusion models while mitigating reward over-optimization. SQDF leverages a KL-regularized RL objective, and introduces a discount factor for better credit assignment, a consistency model for improved Q-function approximation, and an off-policy replay buffer.", "questions": "There are several concerns that may significantly affect the overall evaluation of this paper. \n* In Section 4.1, what is the key difference between the proposed loss (Eq. 11) and the DPOK loss (Eq. 8 in [1])? It seems that the proposed loss allows the use of the reward gradient of the predicted $x_0$ at each timestep. However, the use of reward gradients has already been introduced in DRaFT [2], and the use of rewards on the predicted $x_0$ has been proposed in ReFL [3]. This makes Section 4.1 appear to be a combination of DPOK, DRaFT, and ReFL. \n* In Section 4.2.1, the authors introduce a discount factor to emphasize that later timesteps are more important. However, prior work has also introduced a discount factor but emphasized the earlier denoising steps instead (Eq. 1 in [4]). Could the authors clarify this contradiction? \n* In Section 4.2.2, the authors introduce a consistency model for better $x_0$ prediction。While the $x_0$ predicted by the consistency model may indeed be cleaner, it is no longer the prediction of the original model $p_\\theta$. Why should the reward obtained from this predicted $x_0$ be valid for updating $p_\\theta$? (Even though the images generated by the consistency model and the original model are largely similar, differences always exist in some details, and those details could be crucial for reward evaluation sometimes.) \n* In Section 4.2.3, the authors introduce a replay buffer. In the conclusion, the paper claims that the replay buffer can manage the reward–diversity trade-off. However, according to Appendix E.1, the replay buffer only weights each sample by $\\gamma^t r$, which seems to be equivalent to directly weighting the samples by $\\gamma^t r$ in the loss function. Could the authors provide a more detailed explanation about managing the reward-diversity trade-off? \n\n[1] DPOK: Reinforcement Learning for Fine-tuning Text-to-Image Diffusion Models. \n[2] Directly Fine-Tuning Diffusion Models on Differentiable Rewards. \n[3] ImageReward: Learning and Evaluating Human Preferences for Text-to-Image Generation. \n[4] A Dense Reward View on Aligning Text-to-Image Diffusion with Preference.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* Nice Presentation: The paper addresses reward over-optimization, which is an important issue in diffusion model reinforcement learning. Moreover, the writing of this paper is very easy to understand.\n* Strong Theoretical Grounding: The method is derived from diffusion model MDP formulations, showing solid theoretical consistency. \n* Comprehensive Experiments: Results are thorough, including extensive qualitative and quantitative evaluations, and ablation studies.", "weaknesses": "* Base Model: The paper adopts Stable Diffusion 1.5, an older base model, which raises concerns about the generalizability of the proposed method to more advanced diffusion architectures. \n* Clarity: Some equations and algorithm steps (e.g., Eq. 10-12, Line 14 of Algo.1) could be described more intuitively, as the current notation may make it harder for non-expert readers to follow.", "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)