Paper deep dive
Diffusion Controller: Framework, Algorithms and Parameterization
Tong Yang, Moonkyung Ryu, Chih-Wei Hsu, Guy Tennenholtz, Yuejie Chi, Craig Boutilier, Bo Dai
Intelligence
Status: succeeded | Model: google/gemini-3.1-flash-lite-preview | Prompt: intel-v1 | Confidence: 95%
Last extracted: 3/13/2026, 12:29:31 AM
Summary
The paper introduces 'Diffusion Controller' (DiffCon), a unified control-theoretic framework that models diffusion model fine-tuning as a linearly-solvable Markov Decision Process (LS-MDP). By treating reverse-time diffusion sampling as stochastic control, the authors derive practical reinforcement learning methods, including f-divergence-regularized policy gradients and reward-weighted regression. The framework motivates a 'pretrained + control' model architecture, enabling efficient gray-box adaptation with a frozen backbone, which outperforms existing methods like LoRA in quality-efficiency trade-offs.
Entities (5)
Relation Signals (3)
Diffusion Controller → isbasedon → LS-MDP
confidence 99% · We bridge this gap with Diffusion Controller (DiffCon), a unified control-theoretic view that casts reverse diffusion sampling as state-only stochastic control within (generalized) linearly-solvable Markov Decision Processes (LS-MDPs).
Diffusion Controller → appliedto → Stable Diffusion v1.4
confidence 95% · Experiments on Stable Diffusion v1.4 across supervised and reward-driven finetuning show consistent gains
Diffusion Controller → improvesover → LoRA
confidence 95% · The LS-MDP-motivated score reparameterization delivers a stronger quality–efficiency trade-off—outperforming gray-box baselines and the white-box adapter LoRA.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:Controllable diffusion generation often relies on various heuristics that are seemingly disconnected without a unified understanding. We bridge this gap with Diffusion Controller (DiffCon), a unified control-theoretic view that casts reverse diffusion sampling as state-only stochastic control within (generalized) linearly-solvable Markov Decision Processes (LS-MDPs). Under this framework, control acts by reweighting the pretrained reverse-time transition kernels, balancing terminal objectives against an $f$-divergence cost. From the resulting optimality conditions, we derive practical reinforcement learning methods for diffusion fine-tuning: (i) f-divergence-regularized policy-gradient updates, including a PPO-style rule, and (ii) a regularizer-determined reward-weighted regression objective with a minimizer-preservation guarantee under the Kullback-Leibler (KL) divergence. The LS-MDP framework further implies a principled model form: the optimal score decomposes into a fixed pretrained baseline plus a lightweight control correction, motivating a side-network parameterization conditioned on exposed intermediate denoising outputs, enabling effective gray-box adaptation with a frozen backbone. Experiments on Stable Diffusion v1.4 across supervised and reward-driven finetuning show consistent gains in preference-alignment win rates and improved quality-efficiency trade-offs versus gray-box baselines and even the parameter-efficient white-box adapter LoRA.
Tags
Links
- Source: https://arxiv.org/abs/2603.06981v1
- Canonical: https://arxiv.org/abs/2603.06981v1
Trouble viewing inline? Open PDF directly →
Full Text
132,341 characters extracted from source content.
Expand or collapse full text
Diffusion Controller: Framework, Algorithms and Parameterization Tong Yang CMU & Google Research Carnegie Mellon University; Intern at Google Research; Email: tongyang@andrew.cmu.edu. Moonkyung Ryu Google Research Google; Emails: mkryu,cwhsu,guytenn,cboutilier,bodai@google.com. Chih-Wei Hsu22footnotemark: 2 Google Research Guy Tennenholtz22footnotemark: 2 Google Research Yuejie Chi Yale Yale University; Email: yuejie.chi@yale.edu. Craig Boutilier22footnotemark: 2 Google Research Bo Dai22footnotemark: 2 Google DeepMind (March 2026) Abstract Controllable diffusion generation often relies on various heuristics that are seemingly disconnected without a unified understanding. We bridge this gap with Diffusion Controller (DiffCon), a unified control-theoretic view that casts reverse diffusion sampling as state-only stochastic control within (generalized) linearly-solvable Markov Decision Processes (LS-MDPs). Under this framework, control acts by reweighting the pretrained reverse-time transition kernels, balancing terminal objectives against an f-divergence cost. From the resulting optimality conditions, we derive practical reinforcement learning methods for diffusion fine-tuning: (i) f-divergence-regularized policy-gradient updates, including a PPO-style rule, and (i) a regularizer-determined reward-weighted regression objective with a minimizer-preservation guarantee under the Kullback–Leibler (KL) divergence. The LS-MDP framework further implies a principled model form: the optimal score decomposes into a fixed pretrained baseline plus a lightweight control correction, motivating a side-network parameterization conditioned on exposed intermediate denoising outputs, enabling effective gray-box adaptation with a frozen backbone. Experiments on Stable Diffusion v1.4 across supervised and reward-driven finetuning show consistent gains in preference-alignment win rates and improved quality–efficiency trade-offs versus gray-box baselines and even the parameter-efficient white-box adapter LoRA. Figure 1: Qualitative text-to-image comparisons of pretrained Stable Diffusion v1.4, LoRA finetuning, and ours. Keywords: diffusion, reinforcement learning, controllable generation, text-to-image, representation learning 1 Introduction Recent advances in diffusion models have enabled high-fidelity image synthesis, becoming a default choice for text-to-image generation (Ho et al., 2020; Song and Ermon, 2019), e.g., Stable Diffusion 3 (Esser et al., 2024), Flux.1 Labs et al. (2025), Emu3 (Wang et al., 2024). Yet, controllable generation—steering samples to satisfy user intent, constraints, or downstream objectives—remains challenging, as stronger steering typically requires departing further from the pretrained model and can degrade sample quality (Zhang et al., 2023). In practice, control is achieved through a mix of inference-time mechanisms (e.g., classifier/classifier-free guidance and related guidance variants (Dhariwal and Nichol, 2021; Ho and Salimans, 2022)) and training-time adaptation (e.g., personalization or domain adaptation via paired supervision and parameter-efficient adapters (Ruiz et al., 2023; Zhang et al., 2023)). More recently, a growing line of work fine-tunes diffusion models using human feedback or learned rewards, aiming to optimize terminal objectives such as alignment, aesthetics, or task scores (Lee et al., 2023; Fan et al., 2023). Despite strong empirical progress, these approaches are often presented as a patchwork of objectives and disconnected heuristics, and the field has largely lacked a simple, principled and unified framework to model and analyze the control of these powerful generative processes. 1.1 Our contribution To address this gap, we aim to develop a control-theoretic view of diffusion sampling that provides a common language for interpreting a range of diffusion finetuning approaches. To this end, we view diffusion finetuning as optimal control of the Markov process over intermediate noisy states along the denoising trajectory—induced by the reverse-time sampler. Under this lens, controllable diffusion generation amounts to learning a controlled reverse-time transition kernel of diffusion that steers the induced terminal distribution toward a target objective. Under this controlled-sampling view, diffusion finetuning becomes a state-only stochastic control problem that fits naturally into the linearly-solvable MDP (LS-MDP) framework (Todorov, 2006). In contrast to standard MDP formulations that introduce explicit actions (Uehara et al., 2024a), our LS-MDP–induced diffusion controller (DiffCon) directly modulates the reverse-time dynamics by reweighting the passive (uncontrolled) pretrained transitions. We penalize this modulation with a general f-divergence control cost (recovering the classic KL-regularized LS-MDP (Todorov, 2006) as a special case), which provides a principled trade-off between reward-driven steering and staying close to the pretrained model for preserving stability and sample quality. We next turn the LS-MDP characterization into implementable diffusion finetuning algorithms by realizing the optimality conditions with a score-parameterized diffusion reverse process, rather than directly learning the implied optimal reverse kernels that are hard to sample from. This translation yields practical objectives in the reinforcement learning finetuning (RLFT) setting where no target samples are available and feedback is provided only by a terminal reward model on the final image. Specifically, we derive two RL finetuning updates: i), f-divergence-regularized policy-gradient methods (Williams, 1992) with standard extensions such as PPO (Schulman et al., 2017); and i), regularizer-determined reward-weighted regression objectives with a minimizer-preservation guarantee (under KL, matching the relevant LS-MDP optimal marginal). Finally, beyond offering a unified objective template, the LS-MDP view also sheds light on the optimal model form: the optimal controlled reverse process is a regularized perturbation of the pretrained diffusion model (aka the backbone), so the learned score naturally decomposes into a fixed pretrained baseline plus a small control correction. This motivates our DiffCon parameterization that freezes the backbone and learns a lightweight controller network conditioned on exposed intermediate denoising outputs (e.g., per-step noise predictions or the implied reverse mean), making it naturally compatible with gray-box settings where the backbone architecture is unknown (Poole et al., 2022). Our parameterization steers the backbone model with a small and interpretable module that preserves pretrained stability yet enables strong and tunable control. Our experiments demonstrate the effectiveness of our proposed algorithms and parameterization. Across both SFT and RLFT settings, the resulting objectives significantly improve win rate over the pretrained model. Moreover, the LS-MDP-motivated score reparameterization delivers a stronger quality–efficiency trade-off—outperforming gray-box baselines and the white-box adapter LoRA. 1.2 Related work Reinforcement learning for diffusion models. Existing RL finetuning for diffusion models treats denoising as an MDP and optimizes terminal rewards via policy gradients (e.g., DDPO (Black et al., 2023), DPOK (Fan et al., 2023)), differentiable-reward backpropagation (e.g., DRaFT (Clark et al., 2023)), or reward-weighted likelihood/score matching (Lee et al., 2023; Fan et al., 2023), see Uehara et al. (2024a) for a review. Several works align diffusion models from pairwise preferences using DPO-style objectives (Wallace et al., 2024; Yang et al., 2024a; Kim et al., 2025; Croitoru et al., 2025), and others mitigate delayed supervision by constructing dense or step-level rewards along the denoising trajectory (Yang et al., 2024b). Tang and Zhou (2024) further connects continuous-time diffusion finetuning to entropy/f-divergence regularized stochastic control, with follow-ups casting diffusion finetuning explicitly as entropy-regularized control (Uehara et al., 2024b; Han et al., 2024). Score-function parameterization and modular control. A common diffusion finetuning practice keeps a strong pretrained backbone fixed and steer generation with lightweight adapters (e.g., LoRA (Hu et al., 2022), or DiffFit (Xie et al., 2023), which tunes only bias terms plus a small set of added layerwise scaling factors) or side networks such as ControlNet (Zhang et al., 2023), T2I-Adapter (Mou et al., 2024), IP-Adapter (Ye et al., 2023), Gligen (Li et al., 2023) that add score adjustments (akin to guidance). In contrast, our LS-MDP analysis yields a principled pretrained++controller decomposition of the optimal score that supports gray-box access, with the control term implemented as a side module parameterized via random Fourier features (Rahimi and Recht, 2007; Tancik et al., 2020). Notation. Let ℝ,ℕR,N be the set of real numbers and natural numbers. Let ℝdR^d be the d-dimensional real vector space, and ‖x‖2 \|x \|_2 represent the Euclidean norm of x. The Gaussian distribution with mean μ and covariance matrix Σ is denoted by (x|μ,Σ)N (x|μ, ), and ()U (A ) stands for the uniform distribution over set A. Define [T]≔1,⋯,T[T] \1,·s,T\. Last but not least, (⋅)O(·) and p(⋅)O_p(·) denote the standard asymptotic big-O notation and the big-O in probability notation, respectively. 2 Preliminaries Linearly-Solvable MDP (LS-MDP) (Todorov, 2006). An LS-MDP is a stochastic control problem in which the controller acts directly on the transition kernel: at each time step, it selects a next-state distribution that stays close to a given passive (uncontrolled) dynamic, measured by the Kullback-Leibler (KL) divergence. Concretely, let st∈s_t∈S be the state at time t from the state space S, and ℙt(st+1∣st)P_t(s_t+1 s_t) denote the passive transition kernel at time t. The time-dependent control ut:×→ℝu_t:S×S exponentially tilts the passive dynamic as: ℙu,t(st+1|st)=ℙt(st+1|st)exp(ut(st+1,st)), _u,t (s_t+1|s_t )=P_t (s_t+1|s_t ) (u_t (s_t+1,s_t ) ), s.t.∫ℙu,t(st+1|st)st+1=1,∀t∈[T−1], s.t.\,\, _SP_u,t (s_t+1|s_t )ds_t+1=1, ∀ t∈[T-1], (1) where T is the time horizon. Given an instantaneous reward rt:→ℝr_t:S , the (regularized) value recursion is Vu,t(st)≔rt(st)+st+1∼ℙu,t(⋅|st)[Vu,t+1(st+1)]−τDKL(ℙu,t(⋅|st)∥ℙt(⋅|st)), V_u,t (s_t ) r_t (s_t )+E_s_t+1 _u,t (·|s_t ) [V_u,t+1 (s_t+1 ) ]-τ D_ KL (P_u,t(·|s_t)\,\|\,P_t(·|s_t) ), (2) where DKL(p∥q)≔∫p(x)logp(x)q(x)dxD_ KL(p\,\|\,q) p(x) p(x)q(x)dx is the KL divergence, and τ⩾0τ 0 is the regularization coefficient. The KL-regularized formulation (2) has also been generalized by replacing the KL control cost with more general divergence regularizers, such as Rényi-regularized LS-control (Dvijotham and Todorov, 2011) or Tsallis entropy (Hashizume et al., 2024). Diffusion models (Ho et al., 2020). Diffusion models generate data by learning to iteratively transforms a simple noise distribution into a target data distribution. We describe the setting of conditional generation, where clean images follow the distribution xT∼pdata(⋅∣c)x_T p_ data(· c) conditioned on c∼pcc p_c (e.g., class labels or text), with condition support C.111Our indexing is reversed from the convention used in Ho et al. (2020) to align with common RL notation: xTx_T is the clean image and x1x_1 is pure noise. When c=∅c= , this reduces to unconditional generation. Let βt∈(0,1)t=1T−1\ _t∈(0,1)\_t=1^T-1 be the variance schedule, and define αt=1−βt,α¯t=∏i=tT−1αi. _t=1- _t, α_t= _i=t^T-1 _i. (3) We now describe the forward process and the reverse process used in diffusion models. • Forward process. The forward (noising) process is a fixed Markov chain that gradually corrupts a clean image into (approximately) standard Gaussian noise. Starting from xT∼pdata(⋅∣c)x_T p_ data(· c), for any t∈[T−1]t∈[T-1], xt=αtxt+1+βtξt x_t= _tx_t+1+ _t _t =α¯txT+1−α¯tξ, = α_tx_T+ 1- α_tξ, (4) where ξt,ξ∼(0,d) _t,ξ (0,I_d ). In particular, x1x_1 is approximately standard Gaussian. We use score matching to learn the score function ϵ(⋅,c,t)ε(·,c,t),222To be precise, this is the noise-prediction function, which can be mapped equivalently from the score function (Song et al., 2020). which is parameterized by neural networks as ϵθ(⋅,c,t) _θ(·,c,t), using the following noise-prediction regression objective: minθc,xT,t,ξ[12‖ξ−ϵθ(α¯txT+1−α¯tξ,c,t)‖22], _θE_c,x_T,t,ξ [ 12 \|ξ- _θ ( α_tx_T+ 1- α_tξ,c,t ) \|_2^2 ], where the expectation is taken over c∼pcc p_c, xT∼pdata(⋅|c)x_T p_ data(·|c), t∼([T−1])t ([T-1] ), and ξ∼(0,d)ξ (0,I_d ). • Reverse process. At sampling time, diffusion models iterative denoise a noise sample x1∼(0,d)x_1 (0,I_d ) to xTx_T through the Gaussian reverse-time transitions: xt+1∼pt(⋅|xt,c)≔(⋅|μ(xt,c,t),β~td) x_t+1 p_t (·|x_t,c ) (·|μ(x_t,c,t), β_tI_d) (5) with mean μ(xt,c,t)=1αtxt−βtαt(1−α¯t)ϵθ(xt,c,t)μ(x_t,c,t)= 1 _tx_t- _t _t (1- α_t ) _θ (x_t,c,t ) and reverse-time variance β~t≔βt(1−α¯t+1)1−α¯t β_t _t(1- α_t+1)1- α_t. For unconditional generation, we directly set ϵθ(xt,t)=ϵθ(xt,∅,t) _θ(x_t,t)= _θ(x_t, ,t) to generate the image following the reverse process. For conditional generation, we use the standard classifier-free guidance (CFG) framework (Ho and Salimans, 2022): during training we drop the condition with probability pdrop∈(0,1)p_drop∈(0,1), and at sampling we combine conditional and unconditional predictions as ϵ^θ(xt,c,t)=(1+λCFG)ϵθ(xt,c,t)−λCFGϵθ(xt,∅,t), ε_θ(x_t,c,t)=(1+ _ CFG) _θ(x_t,c,t)- _ CFG _θ(x_t, ,t), where λCFG⩾1 _ CFG 1 controls guidance strength. Access levels for diffusion finetuning. We categorize diffusion fine-tuning by access to the pretrained backbone. In the white-box setting, the backbone can be modified (e.g., full finetuning or inserting adapters such as LoRA (Hu et al., 2022)). In many real-world applications, however, the backbone may be sealed for proprietary or safety reasons, while exposing limited interfaces (e.g., per-timestep noise predictions or denoising trajectories). This gray-box setting motivates finetuning with add-on modules that harvest on these interfaces while keeping the backbone intact (Poole et al., 2022; Yang et al., 2024c; Yu et al., 2023b; Zhu et al., 2025; Yu et al., 2023a). 3 Diffusion Controller through LS-MDP We are interested in finetuning a pretrained diffusion model to generate images from a new distribution (e.g., conditional generation, human-preference alignment, domain adaptation, etc.). In this section, we formulate diffusion finetuning through the lens of controlling LS-MDP, which provides a simpler and unified framework compared to prior works using standard MDPs (Black et al., 2023; Fan et al., 2023)). We call the proposed framework Diffusion Controller (DiffCon): a controllable diffusion reverse process that steers intermediate states to match a target distribution. Building on this formulation, we further derive the corresponding RL finetuning algorithms. 3.1 Framework Prior works (e.g., DDPO (Black et al., 2023), DPOK (Fan et al., 2023)) cast the reverse diffusion chain as a standard MDP by taking the state st=(xt,c,t)s_t=(x_t,c,t) and treating the next sample as the action at=xt+1a_t=x_t+1, with a deterministic environment and a policy given by the reverse transition from xtx_t to xt+1x_t+1. However, this “action” is simply (part of) the next state, so the only real degree of freedom is how we shape the transition kernel itself. This motivates our LS-MDP characterization, which directly controls the reverse transition (regularized to stay close to the pretrained dynamics) rather than introducing a separate action variable, greatly simplifying the presentation. Let ϵ0(xt,c,t) _0(x_t,c,t) be the pretrained score function, whose reverse process is given by x1∼p0,1≔(⋅|0,d)x_1 p_0,1 (·|0,I_d ), and for all t∈[T−1]t∈[T-1]: xt+1∼p0,t(⋅|xt,c)≔(⋅|μ0(xt,c,t),βt~d) x_t+1 p_0,t (·|x_t,c ) (·| _0(x_t,c,t), _tI_d) (6) same as in (5), with mean μ0(xt,c,t)≔1αtxt−βtαt(1−α¯t)ϵ0(xt,c,t). _0 (x_t,c,t ) 1 _tx_t- _t _t (1- α_t ) _0 (x_t,c,t ). (7) We use the generalized LS-MDP framework (c.f. (2), (2)) to characterize the target distribution of finetuning. We define the controllable transition kernels induced by control ut:ℝd×ℝd×→ℝu_t:R^d×R^d×C as ℙu,t(xt+1|xt,c)=p0,t(xt+1|xt,c)exp(ut(xt+1,xt,c)), _u,t(x_t+1|x_t,c)=p_0,t(x_t+1|x_t,c) (u_t(x_t+1,x_t,c)), s.t.∫ℝdℙu,t(xt+1|xt,c)xt+1=1 s.t.\,\, _R^dP_u,t(x_t+1|x_t,c)dx_t+1=1 (8) for all c∈c and t=0,⋯,T−1t=0,·s,T-1. Here we introduce x0≔∅x_0 for controlling the initial distribution p0,0(x1|x0,c)≔p0,1(x1)p_0,0(x_1|x_0,c) p_0,1(x_1) of x1x_1 with u0u_0. Given a reward function rt(xt,c)r_t(x_t,c), we let u⋆=ut⋆t=0T−1u =\u_t \_t=0^T-1 denote the optimal control that solves the following value maximization problem: maxut(⋅,⋅,c)Vu,t(xt,c)≔rt(xt,c)+xt+1∼ℙu,t(⋅|xt,c)[Vu,t+1(xt+1,c)]−τDf(ℙu,t(⋅|xt,c)∥p0,t(⋅|xt,c)) _u_t(·,·,c)V_u,t (x_t,c ) r_t (x_t,c )+E_x_t+1 _u,t (·|x_t,c ) [V_u,t+1 (x_t+1,c ) ]-τ [rgb]0,0,1 [named]pgfstrokecolorrgb0,0,1D_f (P_u,t(·|x_t,c)\,\|\,p_0,t(·|x_t,c) ) (9) for c∈c and t=0,⋯,T−1t=0,·s,T-1 with terminal value function VT(xT,c)≔rT(xT,c)V_T(x_T,c) r_T(x_T,c) and regularization strength τ⩾0τ 0, where DfD_f is an f-divergence (Csiszár, 1963): Df(p∥q)=∫f(p(x)q(x))q(x)xD_f(p\,\|\,q)= f ( p(x)q(x) )q(x)dx for some convex f with f(1)=0f(1)=0. When Df=DKLD_f=D_ KL, (9) recovers the classical KL-regularized LS-MDP (2). Our goal. We aim to learn the optimal transition kernel ℙu⋆,t(⋅|xt,c)P_u ,t(·|x_t,c) that leads to our target distribution pu⋆(⋅|c)p_u (·|c) (given any c∈c )—the terminal sample distribution of xTx_T under the optimal control u⋆u . However, ℙu⋆,t(⋅|xt,c)P_u ,t(·|x_t,c) characterized by (2) is generally intractable to sample from (e.g., in the KL-regularized case, it takes an energy-based form, see (53) in the Appendix), so naively sampling xt+1∼ℙu⋆,t(⋅|xt,c)x_t+1 _u ,t(·|x_t,c) requires expensive inner-loop estimation at each step. To obtain an efficient ancestral sampler, we therefore learn a tractable diffusion reverse process whose induced transitions lead to the terminal distribution pu⋆p_u . Namely, we target for a score function ϵ⋆ε that minimizes the following regression loss ℒorg(ϵ;pu⋆)L_ org(ε;p_u ): 12c,t,ξxT∼pu⋆(⋅|c)‖ξ−ϵ(α¯txT+1−α¯tξ,c,t)‖22, 12E_ subarraycc,t,ξ\\ x_T p_u (·|c) subarray\|ξ-ε ( α_tx_T+ 1- α_tξ,c,t )\|_2^2, (10) where c∼pcc p_c, t∼([T−1])t ([T-1] ), and ξ∼(0,d)ξ (0,I_d ). 3.2 Reinforcement Learning Finetuning (RLFT) When target samples from pu⋆(⋅|c)p_u (·|c) are unavailable and only a reward model is given, we derive tractable RL-style updates from the LS-MDP optimality in (9). Same as existing work (Black et al., 2023; Fan et al., 2023; Lee et al., 2023), we assume a terminal reward model is available: the reward scores only the fully denoised sample at the end of the reverse process, i.e., in (9) we have rt(xt,c)=r(xT,c),t=T,0,otherwise. r_t(x_t,c)= casesr(x_T,c), t=T,\\ 0, otherwise. cases (11) Various diffusion control algorithms can be derived under the LS-MDP framework. To give some examples, we first derive the policy gradient type methods, and then propose reward-weighted losses for RL finetuning. 3.2.1 Policy Gradient and PPO for DiffCon We use pθ,t(xt+1|xt,c)p_θ,t (x_t+1|x_t,c ) to denote the transition kernel with learnable parameters θ. It is often parameterized as the Gaussian (5) for ease of sampling: pθ,t(xt+1|xt,c)=(xt+1|μθ(xt,c,t),β~td) p_θ,t (x_t+1|x_t,c )=N(x_t+1| _θ(x_t,c,t), β_tI_d) (12) with μθ(xt,c,t)≔1αtxt−βtαt(1−α¯t)ϵθ(xt,c,t), _θ(x_t,c,t) 1 _tx_t- _t _t(1- α_t) _θ(x_t,c,t), (13) where ϵθ(xt,c,t) _θ(x_t,c,t) is the score function to be learned. Let pθ,0(x1|x0,c)≔pθ,0(x1|c)p_θ,0 (x_1|x_0,c ) p_θ,0 (x_1|c ), and write xt+1:t′|pθ,xt,c[⋅]E_x_t+1:t |p_θ,x_t,c[·] as the shorthand for xs+1∼pθ,s+1(⋅|xs,c)t⩽s⩽t′−1[⋅]E_ subarraycx_s+1 p_θ,s+1 (·|x_s,c )\\ t s t -1 subarray[·] for any t+1⩽t′⩽Tt+1 t T, and define the value function under pθ,t(⋅|c)p_θ,t (·|c ) as Vθ,t(xt,c) V_θ,t(x_t,c) ≔rt(xt,c)+xt+1|pθ,xt,c[Vθ,t+1(xt+1,c)]−τDf(pθ,t(⋅|xt,c)∥p0,t(⋅|xt,c)). r_t(x_t,c)+E_x_t+1|p_θ,x_t,c [V_θ,t+1(x_t+1,c) ]-τ D_f (p_θ,t(·|x_t,c)\,\|\,p_0,t(·|x_t,c) ). (14) Unfolding (14), under our terminal reward (11), we have Vθ,t(xt,c)=xt+1:T|pθ,xt,c[r(xT,c)−τ∑s=tT−1f(ζθ,s+1)ζθ,s+1], V_θ,t(x_t,c)=E_x_t+1:T|p_θ,x_t,c [r(x_T,c)-τ _s=t^T-1 f( _θ,s+1) _θ,s+1 ], (15) where ζθ,s+1≔pθ,s(xs+1|xs,c)p0,s(xs+1|xs,c) _θ,s+1 p_θ,s (x_s+1|x_s,c )p_0,s (x_s+1|x_s,c ). We define the soft advantage function given (xt,c,t)(x_t,c,t) as Aθ,t+1 A_θ,t+1 ≔Vθ,t+1(xt+1,c)−τf′(ζθ,t+1)−Vθ,t(xt,c)+τxt+1|pθ,xt,c[f′(ζθ,t+1)−f(ζθ,t+1)ζθ,t+1]. V_θ,t+1(x_t+1,c)-τ f ( _θ,t+1 )-V_θ,t(x_t,c)+ _x_t+1|p_θ,x_t,c [f ( _θ,t+1 )- f( _θ,t+1) _θ,t+1 ]. (16) Then Aθ,t+1A_θ,t+1 satisfies xt+1|pθ,xt,c[Aθ,t+1]=0E_x_t+1|p_θ,x_t,c [A_θ,t+1 ]=0. And a special case is when Df=DKLD_f=D_ KL, we have Aθ,t+1=Vθ,t+1(xt+1,c)−τlog(ζθ,s+1)−Vθ,t(xt,c). A_θ,t+1=V_θ,t+1(x_t+1,c)-τ ( _θ,s+1 )-V_θ,t(x_t,c). The following proposition gives the policy gradient expression for maximizing Jθ≔c∼pc,x1∼pθ,1(⋅|c)[Vθ,1(x1,c)]=c∼pc[Vθ,0(x0,c)] J_θ _ subarraycc p_c,\\ x_1 p_θ,1 (·|c ) subarray [V_θ,1(x_1,c) ]=E_c p_c[V_θ,0(x_0,c)] under the generalized LS-MDP framework. Proposition 1 (policy gradient) Under our reward setting (11), the gradient of JθJ_θ is given by ∇θJθ=c∼pc,x1:T|pθ,c[∑t=0T−1∇θlogpθ,t(xt+1|xt,c)Aθ,t+1]. _θJ_θ=E_ subarraycc p_c,\\ x_1:T|p_θ,c subarray [ _t=0^T-1 _θ p_θ,t (x_t+1|x_t,c )A_θ,t+1 ]. The proof of Proposition 1 is given in Appendix B.1. Further, utilizing (16), we also give a PPO update rule for DiffCon with clipping parameter δ∈(0,1)δ∈(0,1): θ(k+1)←θ(k)+η∇θc∼pc,x1:T|pθ(k),c[∑t=0T−1minclip(pθ,t(xt+1|xt,c)pθ(k),t(xt+1|xt,c),1−δ,1+δ)At+1(k), θ^(k+1)←θ^(k)+η _θE_ subarraycc p_c,\\ x_1:T|p_θ^(k),c subarray [ _t=0^T-1 \ clip ( p_θ,t (x_t+1|x_t,c )p_θ^(k),t (x_t+1|x_t,c ),1-δ,1+δ )A_t+1^(k), pθ,t(xt+1|xt,c)pθ(k),t(xt+1|xt,c)At+1(k)]|θ=θ(k). p_θ,t (x_t+1|x_t,c )p_θ^(k),t (x_t+1|x_t,c )A_t+1^(k) \ ] |_θ=θ^(k). (17) Relationship to existing policy gradient methods for diffusion finetuning. Prior RL fine-tuning works such as DDPO (Black et al., 2023) and DPOK (Fan et al., 2023) apply policy gradients to the MDP described at the beginning of Section 3.1, taking the policy to be our reverse transition kernel pθ,t(xt+1|xt,c)p_θ,t(x_t+1|x_t,c). DDPO considers the unregularized setting, and our policy gradient recovers theirs when τ=0τ=0. DPOK adds a KL regularizer toward the pretrained policy, and derives a policy gradient close to ours (but different in the regularization term) when Df=DKLD_f=D_ KL, see Appendix A for a detailed comparison. 3.2.2 Reward-weighted Loss for DiffCon Besides policy-gradient updates, we can obtain ϵ⋆ε by reformulating the intractable objective in (10) into a tractable reward-weighted surrogate that preserves the same minimizer ϵ⋆ε (Ma et al., 2025). To do this, we first introduce the following f-divergence regularized optimization problem on trajectory distribution, where we let P0(⋅|c)P_0(·|c) be the joint reverse trajectory distribution over x1:Tx_1:T under the pretrained model conditioned on c: maxP(⋅|c)P(⋅|c)[r(xT,c)]−τDf(P(⋅|c)∥P0(⋅|c)) _P(·|c)E_P(·|c) [r(x_T,c) ]-τ D_f (P(·|c)\,\|\,P_0(·|c) ) (18) for all c∈c . We let P~⋆(⋅|c) P (·|c) be the optimal solution of (18), and let p~⋆(⋅|c) p (·|c) be the marginal distribution of xTx_T under P~⋆(⋅|c) P (·|c). Theorem 1 gives the tractable reward-weighted loss as the substitute of (10) under KL divergence, by exploiting the equivalence between (18) and (9) when Df=DKLD_f=D_ KL. Moreover, it also gives the reward weighted loss equivalent to ℒorg(⋅;p~⋆)L_ org(·; p ) under general f-divergence regularization. Theorem 1 (reward weighted loss) Under our reward setting (11), when τ>0τ>0, the reward-weighted loss ℒf(ϵ) _f (ε ) ≔12t∼([T−1]),ξ∼(0,d),c∼pc,xT∼p0,T(⋅|c)[wf(r(xT,c),τ)‖ξ−ϵ(α¯txT+1−α¯tξ,c,t)‖22] 12E_ subarrayct ([T-1] ),ξ (0,I_d ),\\ c p_c,x_T p_0,T(·|c) subarray [w_f (r(x_T,c),τ ) \|ξ-ε ( α_tx_T+ 1- α_tξ,c,t ) \|_2^2 ] (19) with the weighting function wf(r(xT,c),τ)=max0,(f′)−1(r(xT,c)−bf.τ(c)τ) w_f (r(x_T,c),τ )= \0,(f )^-1 ( r(x_T,c)-b_f.τ(c)τ ) \ has the same minimum as ℒorg(⋅;p~⋆)L_ org(·; p ), where bf,τ:→ℝb_f,τ:C is a problem-dependent baseline function of condition c. Specially, when Df=DKLD_f=D_ KL, p~⋆=pu⋆ p =p_u . The proof of Theorem 1 is given in Appendix B.2. As special examples of Theorem 1, for Df=DKLD_f=D_ KL, we have the following exponential weighting function wKL(r(xT,c),τ)≔exp(r(xT,c)τ); w_KL(r(x_T,c),τ) ( r(x_T,c)τ ); (20) for Df=DαD_f=D_α (α-divergence), and the following polynomial weighting function wα(r(xT,c),τ)≔[1+α−1τ(r(xT,c)−bα.τ(c))]+1α−1 w_α(r(x_T,c),τ) [1+ α-1τ (r(x_T,c)-b_α.τ(c) ) ]_+ 1α-1 (21) for some baseline function bα,τ:→ℝb_α,τ:C . Note that for general f-divergence Df≠DKLD_f≠ D_ KL, p~⋆ p typically differs from pu⋆p_u , so so there is generally no closed-form weighting wf(r(xT,c),τ)w_f (r(x_T,c),τ ) exactly equivalent to (10). Nevertheless, the general reward-weighted loss form (19) remains a meaningful surrogate—it corresponds to optimizing the score-matching objective under the f-regularized optimal marginal p~⋆ p , and our experiments demonstrate the effectiveness of both the polynomial weighting (21) and the exponential weighting (20). Relationship to existing RWL methods for diffusion finetuning. Prior diffusion RLFT work has also explored reward-weighted objectives: Lee et al. (2023) and Fan et al. (2023, Algorithm 2) optimize a reward-weighted likelihood-style loss on model samples with the weighting being the reward r(xT,c)r(x_T,c) itself. Uehara et al. (2024a, Section 6.1) gives a reward-weighted MLE/score-matching objective with the exponential weighting same as our (20). See Appendix A for a more detailed discussion. 4 Diffusion Controller: Parameterization The LS-MDP formulation implies that the optimal reverse dynamics is not an arbitrary diffusion model—it is the result of a regularized control that stays close to the pretrained reverse kernel (6). This suggests a design principle: rather than fully re-learning a high-capacity score ϵθ _θ, we would like a parameterization that (i), keeps the pretrained score ϵ0 _0 as a strong baseline and isolates the control signal into a lightweight module, and (i) is compatible with gray-box access, where the backbone may be hidden but we can still query intermediate denoising outputs (e.g., ϵ0 _0 or the implied reverse mean μ0 _0). To make this concrete, we focus on the KL-regularized case. Under the assumption that the pretrained score function ϵ0 _0 minimizes the pretrained loss: ℒ0(ϵ)≔12c,t,ξ,xT∼p0(⋅|c)[‖ξ−ϵ(α¯txT+1−α¯tξ,c,t)‖22] _0 (ε ) 12E_ subarraycc,t,ξ,\\ x_T p_0(·|c) subarray [ \|ξ-ε ( α_tx_T+ 1- α_tξ,c,t ) \|_2^2 ] (22) where in the expectation, p0p_0 is the pretraining data distribution, c∼pcc p_c, t∼([T−1])t ([T-1] ), and ξ∼(0,d)ξ (0,I_d ), we next ask: what structure does the LS-MDP optimality impose on the optimal score ϵ⋆ε ? We will introduce Proposition 2 that expresses ϵ⋆ε in terms of ϵ0 _0 plus two control-dependent objects that arise from reweighting the pretrained reverse transition p0,t(⋅|xt,c)p_0,t(·|x_t,c). Further, the specific Gaussian form of p0,t(⋅|xt,c)p_0,t(·|x_t,c) admits a convenient Fourier feature representation that is widely used in representation learning (Rahimi and Recht, 2007; Tancik et al., 2020), allowing us to express the two control-dependent objects as linear combinations of a shared Fourier-style basis function set ϕω(xt,c,t) _ω(x_t,c,t) defined as ϕω(xt,c,t) _ω(x_t,c,t) ≔(cos(ω⊤μ0(xt,c,t)/β~t)sin(ω⊤μ0(xt,c,t)/β~t)). pmatrix (ω _0(x_t,c,t)/ β_t )\\ (ω _0(x_t,c,t)/ β_t ) pmatrix. (23) Now we state the main proposition of this section, where q0,t(⋅|xt,c)q_0,t(·|x_t,c) in the error term is the true posterior transition kernel induced by the data distribution p0p_0 and the forward process (that the learned reverse transition p0,t(⋅|xt,c)p_0,t(·|x_t,c) tries to approximate): Proposition 2 Assume Df=DKLD_f=D_ KL in (9), ϵ0 _0 minimizes the pretrained loss (22), and all xTx_T satisfies ‖xT‖∞⩽1 \|x_T \|_∞ 1. Then the minimum of (10) can be represented as: ϵ⋆(xt,c,t) ε (x_t,c,t) =(1−zt⋆(xt,c))ϵ0(xt,c,t)−1−α¯tβt(zt⋆(xt,c)xt+αt(1−zt⋆(xt,c))ht⋆(xt,c)−αtδ0,t) = (1-z _t(x_t,c) ) _0(x_t,c,t)- 1- α_t _t (z _t(x_t,c)x_t+ _t(1-z _t(x_t,c))h _t(x_t,c)- _t _0,t ) (24) with ht⋆(xt,c) h _t(x_t,c) =1M∑i=1MW⋆(ωi,c,t)ϕωi(xt,c,t)+δh,t, = 1M _i=1^MW ( _i,c,t) _ _i(x_t,c,t)+ _h,t, (25) zt⋆(xt,c)1−zt⋆(xt,c) z _t(x_t,c)1-z _t(x_t,c) =1M∑i=1Mu⋆(ωi,c,t)⊤ϕωi(xt,c,t)+δz,t = 1M _i=1^Mu ( _i,c,t) _ _i(x_t,c,t)+ _z,t (26) for some functions W⋆:ℝd×[T]→ℝd×2W :R^d×C×[T] ^d× 2 and u⋆:ℝd×[T]→ℝ2u :R^d×C×[T] ^2, where ωii=1M∼i.i.d.(⋅|0,d)\ _i\_i=1^M i.i.d. N (·|0,I_d ), M is the sampling number, and the error terms δh,t,δz,t,δ0,t∈ℝd _h,t, _z,t, _0,t ^d satisfy ‖δh,t‖2=p(1M) \| _h,t \|_2=O_p ( 1 M ), ‖δz,t‖2=p(1M) \| _z,t \|_2=O_p ( 1 M ), ‖δ0,t‖2=(βtdTV(q0,t(⋅|xt,c),p0,t(⋅|xt,c))) \| _0,t \|_2=O ( _td TV (q_0,t(·|x_t,c),p_0,t(·|x_t,c) ) ). The proof of Proposition 2 and the hidden problem-dependent constants in O and pO_p are given in Appendix B.3. The main implication of (24) in Proposition 2 indicates the optimal score ϵ⋆ε can be represented as a combination of the pretrained score ϵ0 _0 and two control-dependent objects ht⋆(xt,c)h _t(x_t,c) and zt⋆(xt,c)z _t(x_t,c). Further, (25) and (26) show both h⋆h and the logit of z⋆z can be approximated by a linear combination of shared basis features ϕωi(xt,c,t)i=1M\ _ _i(x_t,c,t)\_i=1^M, where all xtx_t-dependence enters through μ0(xt,c,t) _0(x_t,c,t) in ϕωi(xt,c,t) _ _i(x_t,c,t). This is structurally analogous to cross-attention with μ0(xt,c,t) _0(x_t,c,t) replacing xtx_t as the input, where the condition and time (c,t)(c,t) produce weights that select and mix latent features extracted from μ0 _0. This motivates our DiffCon parameterization of the score function in the following. DiffCon score parameterization. Proposition 2 inspires us to use a “side network” (Zhang et al., 2020) sθ=(zθ1,hθ2)s_θ=(z_ _1,h_ _2) to combine with ϵ0 _0 to parameterize ϵθ _θ as ϵθ(xt,c,t)=ϵ0(xt,c,t)+λmodelsθ(μ0(xt,c,t),c,t) _θ(x_t,c,t)= _0(x_t,c,t)+ _ models_θ( [rgb]0,0,1 [named]pgfstrokecolorrgb0,0,1 _0(x_t,c,t),c,t) (27) with μ0(xt,c,t) _0(x_t,c,t) instead of xtx_t as the input to sθs_θ, where sθ(x,c,t)≔−zθ1(x,c,t)ϵ0(x,c,t)−1−α¯tβt(zθ1(x,c,t)x+αt(1−zθ1(x,c,t))hθ2(x,c,t)). s_θ(x,c,t) -z_ _1(x,c,t) _0(x,c,t)- 1- α_t _t (z_ _1(x,c,t)x+ _t(1-z_ _1(x,c,t))h_ _2(x,c,t) ). (28) In (27), we add a guidance strength λmodel⩾1 _ model 1 on the side network as a hyper-parameter. Motivated by (25) (26), we parameterize zθ1z_ _1 and hθ2h_ _2 with a shared backbone sθs_θ with cross-attention blocks. In practice, we do not explicitly sample ωi _i (which makes the batch size too large); instead we let the feature directions be learnable (implemented implicitly by the cross-attention projections). Importantly, our DiffCon parameterization is orthogonal to the choice of learning algorithms: besides RLFT updates, it can also be trained with standard supervised finetuning (SFT) loss (10) when paired data are available from the target distribution. 5 Experiments 5.1 Setup We test the proposed algorithms and parameterization under three finetuning algorithms for text-to-image generation: 1. SFT: supervised finetuning. The training is identical to CFG, and we conduct it to demonstrate the effectiveness of our parameterization. We use the ”winner images” and their prompts from the Human Preference Dataset (HPD) v2 (Wu et al., 2023) as the training data. We train with a batch size of 128 for 1000 iterations. 2. RWL: finetuning with reward-weighted loss. We optimize over the Human Preference Score (HPS) v2 (Wu et al., 2023) reward model. We use the polynomial reward weighting (21) with τ=τRWL=τ= _ RWL=5e-4, α=1+τRWLα=1+ _ RWL, and set the baseline b as the mean reward of the first training batch. We train with a batch size of 64 for 2000 iterations, sampling online every 2 iterations: 64 prompts generate 64 images with the current score, then we take two gradient steps on the resulting pairs. 3. PPO: finetuning with KL-regularized online PPO. We optimize over the HPS-v2 reward model. We set τ=τKLτ= _ KL (which is the KL regularization coefficient here) to be 1e-4. We train with a batch size of 64 for 2400 iterations. We use stable diffusion v1.4 (Rombach et al., 2022) as the pretrained model. In all methods, we keep the text encoder and VAE fixed, and only finetune the latent score function. For all algorithms, we set the CFG guidance strength λCFG _ CFG to be 7.5 at inference time, and the dropout probability pdropp_ drop to be 0.1. The optimizer is Adam (Kingma and Ba, 2015). To verify the effectiveness of our parameterization, we implement each algorithm with 5 different network structures: • DiffCon (gray-box, ours): we use a side network sθ=(zθ1,hθ2)s_θ=(z_ _1,h_ _2) parameterized and combined with the pretrained score function as specified in (27) with λmodel=1.0 _ model=1.0. sθs_θ consists of 2D-transformer (Vaswani et al., 2017) blocks and ResNet He et al. (2016) blocks, and we take one (linearly transformed) output channel at the final layer to be zθ1z_ _1 (yielding a scalar), and use the remaining channels as hθ2h_ _2. The structure details are in Appendix C. The learning rate under DiffCon is 1e-5 for all algorithms. We initialize the weights of sθs_θ’s last layer to 0 so that, same as LoRA, the initial output is the same as the pretrained model’s output. • DiffCon-Naive (gray-box baseline): we use a side network s¯θ s_θ as in DiffCon that’s almost the same as sθs_θ, except that we do not split the CNN channels at the final layer as is done in DiffCon. Instead, we directly set ϵθ(xt,c,t)=ϵ0(xt,c,t)+λmodels¯θ(xt,c,t). _θ(x_t,c,t)= _0(x_t,c,t)+ _ model s_θ(x_t,c,t). See Appendix C for a more detailed comparison between DiffCon and DiffCon-Naive. The learning rate under DiffCon-Naive is 1e-5 for all algorithms. • LoRA (white-box baseline): LoRA finetuning (Hu et al., 2022) with rank r=16r=16. We set the learning rate under LoRA to be 1e-5 for SFT and 1e-4 for RWL and PPO. • DiffCon-J (white-box, ours): we inject LoRA layers with rank r=4r=4 into the pretrained model and also add DiffCon and jointly train both. For SFT, we set the learning rate for both LoRA and DiffCon to be 1e-5. For RWL and PPO, we set the learning rates for LoRA and DiffCon to be 1e-4 and 1e-5 respectively. • DiffCon-S (white-box, ours): similar to DiffCon-J but we train LoRA and DiffCon separately and combine them at evaluation time. Evaluation. For automatic evaluation of the finetuned models, the main metric is the HPS-v2 win rate (against the pretrained model or the baselines). We also use Clip (Radford et al., 2021), PickScore (Kirstain et al., 2023) and Clip-Aesthetics (Wu et al., 2023) to demonstrate they are not compromised during our finetuning. We found at test time, enlarging the side network guidance strength λmodel _ model can help the finetuned models to achieve better performance, especially for SFT and RWL. Thus we evaluate DiffCon, DiffCon-Naive, DiffCon-J, DiffCon-S with a few λmodel _ model values, and report the best HPS-v2 win rate among them for each checkpoint. We also conduct human evaluation on PPO. See Appendix C for more details. 5.2 Results Main results. Table 1 reports the HPS-v2 win rate improvements against the pretrained model at the end of training for each algorithm, Figure 2 shows the HPS-v2 win rate curves for each algorithm, and Figure 3 shows the win rate comparisons for both our gray-box and white-box methods against their own baselines. They all demonstrate that both our gray-box and white-box methods can surpass their own baselines. Notably, for SFT and RWL, our gray-box method DiffCon (with fewer parameters) can surpass the white-box LoRA, which requires access to the model’s internals. For PPO, our white-box methods can reach a very high win rate (>0.9>0.9) against the pretrained model. Some generation examples can be found in Table 5, 6, 7 in Appendix D. Method # Params ↓ gray-box HPS-v2 win rate vs. pretrained model ↑ SFT (step 1000) RWL (step 2000) PPO (step 2400) Pretrained - - 0.500 0.500 0.500 DiffCon (ours) 1.2×1.2× 10^7 ✓ 0.6667 ± 0.0028 0.6815 ± 0.0111 0.6957 ± 0.0152 DiffCon-Naive 1.2×1071.2× 10^7 ✓ 0.5655 ± 0.0165 0.5060 ± 0.0190 0.5201 ± 0.0203 LoRA 1.7×1071.7× 10^7 0.5766 ± 0.0137 0.6109 ± 0.0059 0.9048 ± 0.0084 DiffCon-J (ours) 1.6×1.6× 10^7 0.6964 ± 0.0293 0.6555 ± 0.0210 0.9353 ± 0.0079 DiffCon-S (ours) 1.6×1.6× 10^7 0.6964 ± 0.0018 0.7091 ± 0.0155 0.9315 ± 0.0164 Table 1: HPS-v2 win rate against the pretrained model on the HPS-v2 test prompt set for all setups (each at its reported checkpoint). Figure 2: Curves of HPS-v2 win rate against the pretrained model for SFT (left), RWL (middle), and PPO (right). (a) HPS-v2 win rate, SFT (b) HPS-v2 win rate, RWL (c) HPS-v2 win rate, PPO (d) human eval., PPO Figure 3: End-of-training win-rate vs. baselines. Each subplot reports three paired comparisons: (gray-box) DiffCon vs. DiffCon-Naive; (white-box) DiffCon-J vs. LoRA; (white-box) DiffCon-S vs. LoRA. (a)-(c): HPS-v2 win rates for SFT/RWL/PPO with orange error bars showing standard deviation; (d): human-evaluated win rate for PPO. More results and ablation studies. Additional results are in Appendix D. Figures 4, 5, 6, and 7 show that increasing the test-time guidance strength λmodel _ model improves generation quality. CLIP, CLIP-Aesthetics, and PickScore sanity checks for the pretrained model and all methods are also provided (Table 8, 9, 10). We also compare DPOK’s reward-weight loss (Fan et al., 2023, Algorithm 2) with our polynomial reward weighting (21) (Figure 8), and ablate key hyperparameters (e.g., τKL _ KL, τRWL _ RWL, reweighting choice (20), learning rates, and side-network architecture). 6 Conclusion We presented DiffCon, a control-theoretic view of reverse diffusion that casts sampling as state-only stochastic control in a (generalized) linearly-solvable MDP. This perspective formalizes control as reweighting the pretrained reverse-time dynamics under an f-divergence regularizer, yielding practical fine-tuning objectives that unify supervised and reward-driven adaptation, including PPO-style updates and reward-weighted regression. DiffCon also motivates a simple pretrained plus controller parameterization: a lightweight side network conditioned on intermediate denoising signals, enabling effective gray-box fine-tuning with the backbone frozen. This improves text-to-image alignment quality–efficiency trade-offs on Stable Diffusion v1.4. A promising direction is to extend the experiments of DiffCon beyond text-to-image alignment to broader diffusion control settings such as personalization, safety alignment, and transfer learning. Acknowledgment The work of Y. Chi is supported in part by NSF grant ECCS-2537078. The work of T. Yang is supported in part by the Wei Shen and Xuehong Zhang Presidential Fellowship at Carnegie Mellon University. The authors thank Krishnamurthy Dvijotham for valuable discussions. References K. Black, M. Janner, Y. Du, I. Kostrikov, and S. Levine (2023) Training diffusion models with reinforcement learning. arXiv preprint arXiv:2305.13301. Cited by: Appendix A, §1.2, §3.1, §3.2.1, §3.2, §3. S. Cen, C. Cheng, Y. Chen, Y. Wei, and Y. Chi (2022) Fast global convergence of natural policy gradient methods with entropy regularization. Operations Research 70 (4), p. 2563–2578. Cited by: §B.1. K. Clark, P. Vicol, K. Swersky, and D. J. Fleet (2023) Directly fine-tuning diffusion models on differentiable rewards. arXiv preprint arXiv:2309.17400. Cited by: §1.2. F. Croitoru, V. Hondru, R. T. Ionescu, N. Sebe, and M. Shah (2025) Curriculum direct preference optimization for diffusion and consistency models. In Proceedings of the Computer Vision and Pattern Recognition Conference, p. 2824–2834. Cited by: §1.2. I. Csiszár (1963) Eine informationstheoretische ungleichung und ihre anwendung auf den beweis der ergodizität von markoffschen ketten. A Magyar Tudományos Akadémia Matematikai Kutató Intézetének Közleményei 8 (1-2), p. 85–108. Cited by: §3.1. P. Dhariwal and A. Nichol (2021) Diffusion models beat gans on image synthesis. Advances in neural information processing systems 34, p. 8780–8794. Cited by: §1. K. Dvijotham and E. Todorov (2011) A unifying framework for linearly solvable control. In UAI 2011, Proceedings of the Twenty-Seventh Conference on Uncertainty in Artificial Intelligence, p. 179–186. Cited by: §2. P. Esser, S. Kulal, A. Blattmann, R. Entezari, J. Müller, H. Saini, Y. Levi, D. Lorenz, A. Sauer, F. Boesel, D. Podell, T. Dockhorn, Z. English, K. Lacey, A. Goodwin, Y. Marek, and R. Rombach (2024) Scaling rectified flow transformers for high-resolution image synthesis. External Links: 2403.03206, Link Cited by: §1. Y. Fan, 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. Advances in Neural Information Processing Systems 36, p. 79858–79885. Cited by: item 2, Appendix A, Appendix A, Appendix D, Appendix D, §1.2, §1, §3.1, §3.2.1, §3.2.2, §3.2, §3, §5.2. Y. Han, M. Razaviyayn, and R. Xu (2024) Stochastic control for fine-tuning diffusion models: optimality, regularity, and convergence. arXiv preprint arXiv:2412.18164. Cited by: §1.2. Y. Hashizume, K. Oishi, and K. Kashima (2024) Tsallis entropy regularization for linearly solvable mdp and linear quadratic regulator. arXiv preprint arXiv:2403.01805. Cited by: §2. K. He, X. Zhang, S. Ren, and J. Sun (2016) Deep residual learning for image recognition. In 2016 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2016, p. 770–778. Cited by: 1st item. J. Ho, A. Jain, and P. Abbeel (2020) Denoising diffusion probabilistic models. Advances in neural information processing systems 33, p. 6840–6851. Cited by: §1, §2, footnote 1. J. Ho and T. Salimans (2022) Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12598. Cited by: §1, §2. E. J. Hu, Y. Shen, P. Wallis, Z. Allen-Zhu, Y. Li, S. Wang, L. Wang, W. Chen, et al. (2022) Lora: low-rank adaptation of large language models.. ICLR 1 (2), p. 3. Cited by: §1.2, §2, 3rd item. D. Kim, S. Lyu, S. W. Kim, and P. H. Seo (2025) Direct diffusion score preference optimization via stepwise contrastive policy-pair supervision. arXiv preprint arXiv:2512.23426. Cited by: §1.2. D. P. Kingma and J. Ba (2015) Adam: A method for stochastic optimization. In 3rd International Conference on Learning Representations, ICLR 2015, Cited by: §5.1. Y. Kirstain, A. Polyak, U. Singer, S. Matiana, J. Penna, and O. Levy (2023) Pick-a-pic: an open dataset of user preferences for text-to-image generation. Advances in neural information processing systems 36, p. 36652–36663. Cited by: §5.1. J. Kwon and V. Perchet (2017) Online Learning and Blackwell Approachability with Partial Monitoring: Optimal Convergence Rates. In Proceedings of the 20th International Conference on Artificial Intelligence and Statistics, A. Singh and J. Zhu (Eds.), Proceedings of Machine Learning Research, Vol. 54, p. 604–613. External Links: Link Cited by: §B.3. B. F. Labs, S. Batifol, A. Blattmann, F. Boesel, S. Consul, C. Diagne, T. Dockhorn, J. English, Z. English, P. Esser, S. Kulal, K. Lacey, Y. Levi, C. Li, D. Lorenz, J. Müller, D. Podell, R. Rombach, H. Saini, A. Sauer, and L. Smith (2025) FLUX.1 kontext: flow matching for in-context image generation and editing in latent space. External Links: 2506.15742, Link Cited by: §1. K. Lee, H. Liu, M. Ryu, O. Watkins, Y. Du, C. Boutilier, P. Abbeel, M. Ghavamzadeh, and S. S. Gu (2023) Aligning text-to-image models using human feedback. arXiv preprint arXiv:2302.12192. Cited by: Appendix A, §1.2, §1, §3.2.2, §3.2. Y. Li, H. Liu, Q. Wu, F. Mu, J. Yang, J. Gao, C. Li, and Y. J. Lee (2023) Gligen: open-set grounded text-to-image generation. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, p. 22511–22521. Cited by: §1.2. H. Ma, T. Chen, K. Wang, N. Li, and B. Dai (2025) Efficient online reinforcement learning for diffusion policy. arXiv preprint arXiv:2502.00361. Cited by: §B.2, §3.2.2. C. Mou, X. Wang, L. Xie, Y. Wu, J. Zhang, Z. Qi, Y. Shan, and X. Qie (2024) T2I-adapter: learning adapters to dig out more controllable ability for text-to-image diffusion models. In Proceedings of the AAAI Conference on Artificial Intelligence (AAAI), External Links: Document Cited by: §1.2. B. Poole, A. Jain, J. T. Barron, and B. Mildenhall (2022) Dreamfusion: text-to-3d using 2d diffusion. arXiv preprint arXiv:2209.14988. Cited by: §1.1, §2. A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, P. Mishkin, J. Clark, et al. (2021) Learning transferable visual models from natural language supervision. In International conference on machine learning, p. 8748–8763. Cited by: §5.1. A. Rahimi and B. Recht (2007) Random features for large-scale kernel machines. Advances in neural information processing systems 20. Cited by: §1.2, §4. R. Rombach, 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, p. 10684–10695. Cited by: §5.1. N. Ruiz, Y. Li, V. Jampani, Y. Pritch, M. Rubinstein, and K. Aberman (2023) Dreambooth: fine tuning text-to-image diffusion models for subject-driven generation. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, p. 22500–22510. Cited by: §1. J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov (2017) Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347. Cited by: §1.1. Y. Song and S. Ermon (2019) Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems 32. Cited by: §1. Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole (2020) Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456. Cited by: footnote 2. M. Tancik, P. Srinivasan, B. Mildenhall, S. Fridovich-Keil, N. Raghavan, U. Singhal, R. Ramamoorthi, J. Barron, and R. Ng (2020) Fourier features let networks learn high frequency functions in low dimensional domains. Advances in neural information processing systems 33, p. 7537–7547. Cited by: §1.2, §4. W. Tang and F. Zhou (2024) Fine-tuning of diffusion models via stochastic control: entropy regularization and beyond. arXiv preprint arXiv:2403.06279. Cited by: §1.2. E. Todorov (2006) Linearly-solvable markov decision problems. Advances in neural information processing systems 19. Cited by: §1.1, §2. M. Uehara, Y. Zhao, T. Biancalani, and S. Levine (2024a) Understanding reinforcement learning-based fine-tuning of diffusion models: a tutorial and review. arXiv preprint arXiv:2407.13734. Cited by: Appendix A, §1.1, §1.2, §3.2.2. M. Uehara, Y. Zhao, K. Black, E. Hajiramezanali, G. Scalia, N. L. Diamant, A. M. Tseng, T. Biancalani, and S. Levine (2024b) Fine-tuning of continuous-time diffusion models as entropy-regularized control. arXiv preprint arXiv:2402.15194. Cited by: §1.2. A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin (2017) Attention is all you need. In Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, p. 5998–6008. Cited by: 1st item. P. von Platen, S. Patil, A. Lozhkov, P. Cuenca, N. Lambert, K. Rasul, M. Davaadorj, and T. Wolf (2022) Diffusers: state-of-the-art diffusion models. GitHub. Note: https://github.com/huggingface/diffusersVersion X.Y.Z Cited by: Appendix C. B. Wallace, M. Dang, R. Rafailov, L. Zhou, A. Lou, S. Purushwalkam, S. Ermon, C. Xiong, S. Joty, and N. Naik (2024) Diffusion model alignment using direct preference optimization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, p. 8228–8238. Cited by: §1.2. X. Wang, X. Zhang, Z. Luo, Q. Sun, Y. Cui, J. Wang, F. Zhang, Y. Wang, Z. Li, Q. Yu, Y. Zhao, Y. Ao, X. Min, T. Li, B. Wu, B. Zhao, B. Zhang, L. Wang, G. Liu, Z. He, X. Yang, J. Liu, Y. Lin, T. Huang, and Z. Wang (2024) Emu3: next-token prediction is all you need. External Links: 2409.18869, Link Cited by: §1. R. Williams (1992) Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine Learning 8 (3-4), p. 229–256. Cited by: §1.1. X. Wu, Y. Hao, K. Sun, Y. Chen, F. Zhu, R. Zhao, and H. Li (2023) Human preference score v2: a solid benchmark for evaluating human preferences of text-to-image synthesis. arXiv preprint arXiv:2306.09341. Cited by: item 1, item 2, §5.1. E. Xie, L. Yao, H. Shi, Z. Liu, D. Zhou, Z. Liu, J. Li, and Z. Li (2023) Difffit: unlocking transferability of large diffusion models via simple parameter-efficient fine-tuning. In Proceedings of the IEEE/CVF International Conference on Computer Vision, p. 4230–4239. Cited by: §1.2. K. Yang, J. Tao, J. Lyu, C. Ge, J. Chen, W. Shen, X. Zhu, and X. Li (2024a) Using human feedback to fine-tune diffusion models without any reward model. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, p. 8941–8951. Cited by: §1.2. S. Yang, T. Chen, and M. Zhou (2024b) A dense reward view on aligning text-to-image diffusion with preference. arXiv preprint arXiv:2402.08265. Cited by: §1.2. S. Yang, Y. Du, B. Dai, D. Schuurmans, J. B. Tenenbaum, and P. Abbeel (2024c) Probabilistic adaptation of black-box text-to-video models. In The Twelfth International Conference on Learning Representations, Cited by: §2. H. Ye, J. Zhang, S. Liu, X. Han, and W. Yang (2023) Ip-adapter: text compatible image prompt adapter for text-to-image diffusion models. arXiv preprint arXiv:2308.06721. Cited by: §1.2. L. Yu, W. Xiang, and K. Han (2023a) Edit-diffnerf: editing 3d neural radiance fields using 2d diffusion model. arXiv preprint arXiv:2306.09551. Cited by: §2. X. Yu, Y. Guo, Y. Li, D. Liang, S. Zhang, and X. Qi (2023b) Text-to-3d with classifier score distillation. arXiv preprint arXiv:2310.19415. Cited by: §2. J. O. Zhang, A. Sax, A. Zamir, L. J. Guibas, and J. Malik (2020) Side-tuning: A baseline for network adaptation via additive side networks. In Computer Vision - ECCV 2020 - 16th European Conference, Glasgow, Lecture Notes in Computer Science, Vol. 12348, p. 698–714. Cited by: §4. L. Zhang, A. Rao, and M. Agrawala (2023) Adding conditional control to text-to-image diffusion models. arXiv preprint arXiv:2302.05543. Cited by: §1.2, §1. H. Zhu, Y. Xu, C. Xu, T. Shen, W. Liu, Y. Du, J. Yu, and S. He (2025) Stable score distillation. In Proceedings of the IEEE/CVF International Conference on Computer Vision, p. 16597–16606. Cited by: §2. Appendix A Additional Details on RLFT Algorithms Relationship to existing RWL methods for diffusion finetuning. Lee et al. [2023] proposes a heuristic reward-weighted loss for finetuning the diffusion model: they minimize a reward-weighted negative log-likelihood (NLL) on the generated data weighted by reward r(xT,c)r(x_T,c) itself. Fan et al. [2023, Algorithm 2] also uses a loss that is an upper bound of the reward-weighted NLL with the same linear weighting. Uehara et al. [2024a, Section 6.1] proposes a reward-weighted MLE algorithm with the following loss (under our notation): c,t,ξx1:t−1∼pθ,xt:T∼p0[exp(r(xT,c)τ)‖βtαt(1−α¯t)(ϵθ(xt,c,t)−ϵ0(xt,c,t))+ξ‖22], _ subarraycc,t,ξ\\ x_1:t-1 p_θ,\\ x_t:T p_0 subarray [ ( r (x_T,c )τ ) \| _t _t(1- α_t)( _θ(x_t,c,t)- _0(x_t,c,t))+ξ \|_2^2 ], (29) where c∼pcc p_c, t∼([T−1])t ([T-1] ), ξ∼(0,d)ξ (0,I_d ), ϵθ _θ is their (unconditional) score function, and we use pθp_θ to denote the transition kernel under the current parameter θ. We can see that their score function could be seen as a (linear) reparametrization of our (unconditional) score function ϵθ⋆ε _θ. And they use a mixture of online and offline sampling. Relationship to existing policy gradient-type methods for diffusion finetuning. Prior RL fine-tuning works such as DDPO Black et al. [2023] and DPOK Fan et al. [2023] propose policy gradient methods under the MDP framework described at the end of Section 3.1, which takes the policy to be the reverse transition kernel pθ,t(xt+1∣xt,c)p_θ,t(x_t+1 x_t,c). Converting to our notation, the following relationships hold: 1. When τ=0τ=0, our policy-gradient reduces to the standard REINFORCE gradient used in DDPO: ∇θDDPO _θJ_θ DDPO =c∼pc,x1:T∼pθ(⋅|c)[∑t=0T−1∇θlogpθ,t(xt+1∣xt,c)r(xT,c)] =E_ subarraycc p_c,\\ x_1:T p_θ(·\,|\,c) subarray [ _t=0^T-1 _θ p_θ,t(x_t+1 x_t,c)\;r(x_T,c) ] (30) up to the usual addition of baselines (which do not change the expectation of the gradient). 2. DPOK Fan et al. [2023] also incorporates a regularization term toward the pretrained model. However, DPOK optimizes a different surrogate objective (based on tractable approximations/upper bounds for the KL-to-reference term), which leads to a policy gradient with a different regularization term than ours under KL divergence: ∇θDPOK=c∼pc,x1:T∼pθ(⋅|c)[∑t=0T−1∇θlogpθ,t(xt+1∣xt,c)(r(xT,c)−τlogpθ,t(xt+1∣xt,c)p0,t(xt+1∣xt,c))]. _θJ_θ DPOK=E_ subarraycc p_c,\\ x_1:T p_θ(·|c) subarray [ _t=0^T-1 _θ p_θ,t(x_t+1 x_t,c)\; (r(x_T,c)- [rgb]0,0,1 [named]pgfstrokecolorrgb0,0,1τ p_θ,t(x_t+1 x_t,c)p_0,t(x_t+1 x_t,c) ) ]. (31) In comparison, by (15) we can see that when Df=DKLD_f=D_ KL, our regularization term at each step t is a sum τ∑s=tT−1logpθ,s(xs+1∣xs,c)p0,s(xs+1∣xs,c). τ _s=t^T-1 p_θ,s(x_s+1 x_s,c)p_0,s(x_s+1 x_s,c). (32) Appendix B Missing Proofs Notation. We usually use q to denote the distributions induced by the forward process, and use p to denote the distributions induced by the reverse process. B.1 Proof of Proposition 1 This proof follows a similar idea to the proof of Lemma 6 in Cen et al. [2022]. We let Dθ,t≔Df(pθ,t(⋅|xt,c)∥p0,t(⋅|xt,c))D_θ,t D_f (p_θ,t(·|x_t,c)\,\|\,p_0,t(·|x_t,c) ) for notation simplicity. By (14), we have for any c∈c : ∇θVθ,0(x0,c) _θV_θ,0(x_0,c) =∇θ[∫x1pθ,0(x1|x0,c)(r0(x0,c)+Vθ,1(x1,c))x1]−τ∇θDθ,0 = _θ [ _x_1p_θ,0 (x_1|x_0,c ) (r_0(x_0,c)+V_θ,1(x_1,c) )dx_1 ]-τ _θD_θ,0 =∫x1(pθ,0(x1|x0,c)∇θlogpθ,0(x1|x0,c))(r0(x0,c)+Vθ,1(x1,c))x1−τ∇θDθ,0 = _x_1 (p_θ,0 (x_1|x_0,c ) _θ p_θ,0 (x_1|x_0,c ) ) (r_0(x_0,c)+V_θ,1(x_1,c) )dx_1-τ _θD_θ,0 +∫x1pθ,0(x1|x0,c)∇θVθ,1(x1,c)x1. + _x_1p_θ,0 (x_1|x_0,c ) _θV_θ,1(x_1,c)dx_1. (33) Note that for all t=0,1,⋯,T−1t=0,1,·s,T-1, we have ∫xt+1pθ,t(xt+1|xt,c)∇θlogpθ,t(xt+1|xt,c)xt+1=∇θ∫xt+1pθ,t(xt+1|xt,c)xt+1=∇θ1=0, _x_t+1p_θ,t (x_t+1|x_t,c ) _θ p_θ,t (x_t+1|x_t,c )dx_t+1= _θ _x_t+1p_θ,t (x_t+1|x_t,c )dx_t+1= _θ1=0, (34) and (recall we define ζθ,s+1≔pθ,s(xs+1|xs,c)p0,s(xs+1|xs,c) _θ,s+1 p_θ,s (x_s+1|x_s,c )p_0,s (x_s+1|x_s,c )) ∇θDθ,t _θD_θ,t =∫f′(ζθ,t+1)∇θpθ,t(xt+1|xt,c)xt+1 = f ( _θ,t+1 ) _θp_θ,t (x_t+1|x_t,c )dx_t+1 =∫pθ,t(xt+1|xt,c)f′(ζθ,t+1)∇θlogpθ,t(xt+1|xt,c)xt+1. = p_θ,t (x_t+1|x_t,c )f ( _θ,t+1 ) _θ p_θ,t (x_t+1|x_t,c )dx_t+1. (35) Letting t=0t=0 in (34), (B.1), and plugging them into (B.1), we have ∇θVθ,0(x0,c) _θV_θ,0(x_0,c) =x1∼pθ,0(⋅|x0,c)[∇θlogpθ,0(x1|x0,c)(−τf′(ζθ,1)+Vθ,1(x1,c))+∇θVθ,1(x1,c)], =E_x_1 p_θ,0 (·|x_0,c ) [ _θ p_θ,0 (x_1|x_0,c ) (-τ f ( _θ,1 )+V_θ,1(x_1,c) )+ _θV_θ,1(x_1,c) ], (36) where pθ,0(x1|x0,c)=pθ,0(x1|c)p_θ,0 (x_1|x_0,c )=p_θ,0 (x_1|c ). Repeating the above process for t=1,…,T−1t=1,…,T-1, and taking expectation over c∼pcc p_c, we obtain the following policy gradient: ∇θJθ _θJ_θ =∇θc∼pc[Vθ,1(x1,c)] = _θE_c p_c [V_θ,1(x_1,c) ] =c∼pc,x1:T|pθ,c[∑t=0T−1(∇θlogpθ,t(xt+1|xt,c))(Vθ,t+1(xt+1,c)−τf′(ζθ,t+1))⏟≔Qθ,t+1], =E_ subarraycc p_c,\\ x_1:T|p_θ,c subarray [ _t=0^T-1 ( _θ p_θ,t (x_t+1|x_t,c ) ) (V_θ,t+1(x_t+1,c)-τ f ( _θ,t+1 ) )_ Q_θ,t+1 ], (37) where we define the soft Q-function Qθ,t+1≔Vθ,t+1(xt+1,c)−τf′(ζθ,t+1). Q_θ,t+1 V_θ,t+1(x_t+1,c)-τ f ( _θ,t+1 ). (38) Note that Aθ,t+1 A_θ,t+1 =(16)Vθ,t+1(xt+1,c)−τf′(ζθ,t+1)−Vθ,t(xt,c)−τDf(pθ,t(⋅|xt,c)∥p0,t(⋅|xt,c)) eq:advantage=V_θ,t+1(x_t+1,c)-τ f ( _θ,t+1 )-V_θ,t(x_t,c)-τ D_f (p_θ,t(·|x_t,c)\,\|\,p_0,t(·|x_t,c) ) +τxt+1∼pθ,t(⋅|xt,c)[f′(ζθ,t+1)] + _x_t+1 p_θ,t (·|x_t,c ) [f ( _θ,t+1 ) ] =(14)Vθ,t+1(xt+1,c)−τf′(ζθ,t+1)−(xt+1∼pθ,t(⋅|xt,c)[Vθ,t+1(xt+1,c)]−τxt+1∼pθ,t(⋅|xt,c)[f′(ζθ,t+1)]) eq:V_theta=V_θ,t+1(x_t+1,c)-τ f ( _θ,t+1 )-(E_x_t+1 p_θ,t (·|x_t,c ) [V_θ,t+1(x_t+1,c) ]- _x_t+1 p_θ,t (·|x_t,c ) [f ( _θ,t+1 ) ]) =(38)Qθ,t+1−xt+1∼pθ,t(⋅|xt,c)[Qθ,t+1], eq:soft_Q=Q_θ,t+1-E_x_t+1 p_θ,t (·|x_t,c ) [Q_θ,t+1 ], (39) where the second equality also uses the fact that rt(xt,c)=0r_t(x_t,c)=0 (c.f. (11)). Since the term xt+1∼pθ,t(⋅|xt,c)[Qθ,t+1]E_x_t+1 p_θ,t (·|x_t,c ) [Q_θ,t+1 ] in (B.1) is a constant to xt+1x_t+1, by (B.1) and (34), we have ∇θJθ=c∼pc,x1:T|pθ,c[∑t=0T−1(∇θlogpθ,t(xt+1|xt,c))Aθ,t+1]. _θJ_θ=E_ subarraycc p_c,\\ x_1:T|p_θ,c subarray [ _t=0^T-1 ( _θ p_θ,t (x_t+1|x_t,c ) )A_θ,t+1 ]. (40) B.2 Proof of Theorem 1 Step 1: Prove the equivalence between the reward-weighted loss (19) and ℒorg(⋅;p~⋆)L_ org(·; p ) under f-divergence. We’l utilize the following lemma in this step, which is adapted from Proposition 3.1 in Ma et al. [2025]: Lemma 1 Given any target distribution p⋆(⋅|c)p (·|c) of xTx_T, let qt⋆(⋅|c)q _t(·|c) be the forward marginal distribution of xt≔α¯txT+1−α¯tξx_t α_tx_T+ 1- α_tξ with ξ∼(0,d)ξ (0,I_d ), xT∼p⋆(⋅|c)x_T p (·|c). Then for any g:ℝd×[T]→(0,+∞)g:R^d×C×[T]→(0,+∞), ϵε and c∈c , we define ℒ(ϵ)≔12c∼pc,t∼([T−1])[∫g(xt,c,t)∥ϵ(xt,c,t)+1−α¯t∇xtlogqt⋆(xt|c)∥22dxt]. (ε ) 12E_c p_c,t ([T-1] ) [ g (x_t,c,t ) \|ε (x_t,c,t )+ 1- α_t _x_t q _t(x_t|c) \|_2^2dx_t ]. (41) Then for any g:ℝd×[T]→(0,+∞)g:R^d×C×[T]→(0,+∞) that makes (41) well-defined, ℒorg(⋅;p⋆)L_ org(·;p ) and ℒL has the same minimum ϵp⋆ _p , which satisfies ∀(xt,c,t)∈ℝd×[T]:ϵp⋆(xt,c,t)=−1−α¯t∇xtlogqt⋆(xt|c). ∀(x_t,c,t) ^d×C×[T]: _p (x_t,c,t )=- 1- α_t _x_t q _t(x_t|c). (42) The proof of Lemma 1 is given in Appendix B.4. We define the density ratio w as w(x1:T,c)≔P(x1:T|c)P0(x1:T|c). w(x_1:T,c) P (x_1:T|c )P_0(x_1:T|c). (43) Then we can rewrite the optimization problem (18) as maxw(⋅,c) _w(·,c) x1:T∼P0(⋅|c)[r(xT,c)w(x1:T,c)−τf(w(x1:T,c))] _x_1:T P_0(·|c) [r(x_T,c)w(x_1:T,c)-τ f (w(x_1:T,c) ) ] s.t.w(x1:T,c)⩾0,x1:T∼P0(⋅|c)[w(x1:T,c)]=1. s.t. w(x_1:T,c) 0, _x_1:T P_0(·|c) [w(x_1:T,c) ]=1. (44) We let w⋆w be the optimal solution of (B.2). Introduce a multiplier b(c)∈ℝb(c) for the constraint x1:T∼P0(⋅|c)[w(x1:T,c)]=1E_x_1:T P_0(·|c) [w(x_1:T,c) ]=1. The Lagrangian is L(w(x1:T,c),b(c))=x1:T∼P0(⋅|c)[r(xT,c)w(x1:T,c)−τf(w(x1:T,c))−b(c)w(x1:T,c)]+b(c). L(w(x_1:T,c),b(c))=E_x_1:T P_0(·|c) [r(x_T,c)w(x_1:T,c)-τ f (w(x_1:T,c) )-b(c)w(x_1:T,c) ]+b(c). (45) We consider Gateaux derivative and take w↦w+ϵhw w+ε h with arbitrary bounded h. Differentiating at ϵ=0ε=0, we have dϵL(w+ϵh,b)|ϵ=0=∫r(xT,c)−τf′(w(x1:T,c))−b(c)⏟=:ϕ(x1:T,c)h(x1:T)P0(x1:T|c)x1:T. ddεL(w+ε h,b) |_ε=0= \r(x_T,c)-τ f (w(x_1:T,c) )-b(c) \_=:φ(x_1:T,c)h(x_1:T)P_0(x_1:T|c)dx_1:T. (46) By KKT condition, if w⋆(x1:T,c)>0w (x_1:T,c)>0, ϕ(x1:T,c)=0φ(x_1:T,c)=0; if w⋆(x1:T,c)=0w (x_1:T,c)=0, ϕ(x1:T,c)⩽0φ(x_1:T,c) 0 (no profitable increase). Therefore, we have w⋆(x1:T,c)=[(f′)−1(r(xT,c)−bf,τ(c)τ)]+, w (x_1:T,c)= [(f )^-1 ( r(x_T,c)-b_f,τ(c)τ ) ]_+, (47) for some bf,τ(c)∈ℝb_f,τ(c) , where [x]+≔max(0,x)[x]_+ (0,x ), and bf,τ(c)b_f,τ(c) is chosen so that x1:T∼P0(⋅|c)[(f′)−1(r(xT,c)−bf,τ(c)τ)]+=1.E_x_1:T P_0(·|c) [(f )^-1 ( r(x_T,c)-b_f,τ(c)τ ) ]_+=1. This suggests p~⋆(xT|c) p (x_T|c) satisfies p~⋆(xT|c) p (x_T|c) =∫P~⋆(x1:T|c)x1:T−1 = P (x_1:T|c)dx_1:T-1 =∫P0(x1:T|c)x1:T−1[(f′)−1(r(xT,c)−bf,τ(c)τ)]+ = P_0(x_1:T|c)dx_1:T-1 [(f )^-1 ( r(x_T,c)-b_f,τ(c)τ ) ]_+ =p0,T(xT|c)[(f′)−1(r(xT,c)−bf,τ(c)τ)]+, =p_0,T(x_T|c) [(f )^-1 ( r(x_T,c)-b_f,τ(c)τ ) ]_+, (48) where p0,T(xT|c)p_0,T(x_T|c) is the reverse marginal distribution of xTx_T under the pretrained model. Let g(xt,c,t)=q~t⋆(xt|c) g(x_t,c,t)= q_t (x_t|c ) (49) in Lemma 1, where q~t⋆(xt|c) q_t (x_t|c ) is the forward marginal distribution of xt=α¯txT+1−α¯tξx_t= α_tx_T+ 1- α_tξ with ξ∼(0,d)ξ (0,I_d ) and xT∼p~⋆(⋅|c)x_T p (·|c ). Then we have ∫g(xt,c,t)∥ϵ(xt,c,t)+1−α¯t∇xtlogq~t⋆(xt|c)∥22dxt g (x_t,c,t ) \|ε (x_t,c,t )+ 1- α_t _x_t q_t (x_t|c ) \|_2^2dx_t =∫g(xt,c,t)q~t⋆(xt|c)∫p~⋆(xT|c)qt(xt|xT)∥ϵ(xt,c,t)+1−α¯t∇xtlogqt(xt|xT)∥22dxTdxt+ = g (x_t,c,t ) q _t(x_t|c) p (x_T|c )q_t(x_t|x_T) \|ε (x_t,c,t )+ 1- α_t _x_t q_t(x_t|x_T) \|_2^2dx_Tdx_t+ const =∫p0,T(xT|c)[(f′)−1(r(xT,c)−bf,τ(c)τ)]+qt(xt|xT)∥ϵ(xt,c,t)+1−α¯t∇xtlogqt(xt|xT)∥22dxTdxt+ = p_0,T(x_T|c) [(f )^-1 ( r(x_T,c)-b_f,τ(c)τ ) ]_+q_t(x_t|x_T) \|ε (x_t,c,t )+ 1- α_t _x_t q_t(x_t|x_T) \|_2^2dx_Tdx_t+ const =xT∼p0,T,ξ∼(0,d)[[(f′)−1(r(xT,c)−bf,τ(c)τ)]+‖ξ−ϵ(α¯txT+1−α¯tξ,c,t)‖22]+, =E_x_T p_0,T, ξ (0,I_d ) [ [(f )^-1 ( r(x_T,c)-b_f,τ(c)τ ) ]_+ \|ξ-ε ( α_tx_T+ 1- α_tξ,c,t ) \|_2^2 ]+ const, (50) where the second line is obtained by replacing q⋆(xt|c)q (x_t|c) and p⋆(xT|c)p (x_T|c) in (B.4) in the proof of Lemma 1 by q~t⋆(xt|c) q_t (x_t|c ) and p~⋆(xT|c) p (x_T|c ), the third line is by our choice of g in (49) and (B.2). Therefore, by Lemma 1, under (49), ℒ(ϵ)L(ε) has the same minimum as ℒorg(⋅;p~⋆)L_ org(·; p ). Step 2: show p~⋆(xT|c)=pu⋆(xT|c) p (x_T|c)=p_u (x_T|c) when Df=DKLD_f=D_ KL. When Df=DKLD_f=D_ KL, the optimal KL-regularized Bellman equation is (t=0,⋯,T−1t=0,·s,T-1) ∀c∈:Vt⋆(xt,c) ∀ c : V _t (x_t,c ) =maxu(⋅,xt,c)rt(xt,c)+xt+1∼ℙu,t(⋅|xt,c)[Vt+1⋆(xt+1,c)−τut(xt+1,xt,c)] = _u (·,x_t,c )\,\,r_t (x_t,c )+E_x_t+1 _u,t (·|x_t,c ) [V _t+1 (x_t+1,c )-τ u_t (x_t+1,x_t,c ) ] =maxu(⋅,xt,c)Vu,t(xt,c), = _u (·,x_t,c )V_u,t (x_t,c ), (51) which induces ∀c∈:ut⋆(xt+1,xt,c) ∀ c : u _t (x_t+1,x_t,c ) ≔argmaxu(⋅,xt,c)Vu,t(xt,c) _u (·,x_t,c )V_u,t (x_t,c ) =1τVt+1⋆(xt+1,c)−log(xt+1∼p0,t(⋅|xt,c)[exp(1τVt+1⋆(xt+1,c))]). = 1τV _t+1 (x_t+1,c )- (E_x_t+1 p_0,t (·|x_t,c ) [ ( 1τV _t+1 (x_t+1,c ) ) ] ). (52) Then by (3.1), the optimal transition operator becomes ℙu⋆,t(xt+1|xt,c)=p0,t(xt+1|xt,c)exp(1τVt+1⋆(xt+1,c))xt+1∼p0,t(⋅|xt,c)[exp(1τVt+1⋆(xt+1,c))]. _u ,t (x_t+1|x_t,c )=p_0,t (x_t+1|x_t,c ) ( 1τV _t+1 (x_t+1,c ) )E_x_t+1 p_0,t (·|x_t,c ) [ ( 1τV _t+1 (x_t+1,c ) ) ]. (53) In addition, plugging (B.2) into (B.2), we have Vt⋆(xt,c)=rt(xt,c)+τlog(xt+1∼p0,t(⋅|xt,c)[exp(1τVt+1⋆(xt+1,c))]). V _t (x_t,c )=r_t (x_t,c )+τ (E_x_t+1 p_0,t (·|x_t,c ) [ ( 1τV _t+1 (x_t+1,c ) ) ] ). (54) Define Zt(xt,c)≔exp(1τVt⋆(xt,c)), Z_t (x_t,c ) ( 1τV _t (x_t,c ) ), (55) then under our terminal reward setting (11), we have ZT(xT,c)≔exp(1τr(xT,c)), Z_T (x_T,c ) ( 1τr (x_T,c ) ), (56) and the following recursion given by (54): ∀t=0,⋯,T−1:Zt(xt,c)=xt+1∼p0,t(⋅|xt,c)[Zt+1(xt+1,c)]. ∀ t=0,·s,T-1: Z_t (x_t,c )=E_x_t+1 p_0,t (·|x_t,c ) [Z_t+1 (x_t+1,c ) ]. (57) Especially, when t=0t=0, we have (recall we define x0≔∅x_0 ) Z0(c)≔Z0(x0,c)=x1∼(0,d)[Z1(x1,c)]. Z_0(c) Z_0(x_0,c)=E_x_1 (0,I_d ) [Z_1 (x_1,c ) ]. (58) Plugging (57) into (53), we have for all t=0,⋯,T−1t=0,·s,T-1: ℙu⋆,t(xt+1|xt,c)=p0,t(xt+1|xt,c)Zt+1(xt+1,c)Zt(xt,c). _u ,t (x_t+1|x_t,c )=p_0,t (x_t+1|x_t,c ) Z_t+1 (x_t+1,c )Z_t (x_t,c ). (59) Let ℙu⋆(x1:T|c)P_u (x_1:T|c) be the joint distribution of x1:Tx_1:T under the optimal control u⋆u . Then by (59), we have ℙu⋆(x1:T|c) _u (x_1:T|c) =pu⋆,−1(x0|c)∏t=0T−1ℙu⋆,t(xt+1|xt,c) =p_u ,-1(x_0|c) _t=0^T-1P_u ,t (x_t+1|x_t,c ) =∏t=0T−1p0,t(xt+1|xt,c)Zt+1(xt+1,c)Zt(xt,c) = _t=0^T-1p_0,t (x_t+1|x_t,c ) Z_t+1 (x_t+1,c )Z_t (x_t,c ) =P0(x1:T|c)ZT(xT,c)Z0(c) =P_0(x_1:T|c) Z_T (x_T,c )Z_0(c) =(56)P0(x1:T|c)exp(1τr(xT,c))Z0(c). eq:Z_T=P_0(x_1:T|c) ( 1τr (x_T,c ) )Z_0(c). (60) Integrating out x1:T−1x_1:T-1, we have pu⋆(xT|c)=p0,T(xT|c)exp(1τr(xT,c))Z0(c), p_u (x_T|c)=p_0,T(x_T|c) ( 1τr (x_T,c ) )Z_0(c), (61) where p0,T(xT|c)p_0,T(x_T|c) is the reverse marginal distribution of xTx_T under the pretrained model. Note KL divergence is f-divergence with f(t)=tlogtf(t)=t t. By (B.2), we have p~⋆(xT|c)∝p0,T(xT|c)exp(1τr(xT,c)), p (x_T|c) p_0,T(x_T|c) ( 1τr (x_T,c ) ), (62) where we use the fact that (f′)−1(y)=exp(y−1)(f )^-1(y)= (y-1) for KL divergence. This gives the desired result p~⋆(xT|c)=pu⋆(xT|c) p (x_T|c)=p_u (x_T|c). B.3 Proof of Proposition 2 Let q0,t(xt|c),qu⋆,t(xt|c)q_0,t(x_t|c),q_u ,t(x_t|c) denote the marginal distributions of xtx_t induced by the forward process starting from target distributions p0p_0 and pu⋆p_u , respectively, i.e., q0,t(⋅|c)q_0,t(·|c) (resp. qu⋆,t(⋅|c)q_u ,t(·|c)) is the probability of xt=α¯txT+1−α¯tξtx_t= α_tx_T+ 1- α_t _t, where ξt∼(0,d) _t (0,I_d ), xT∼p0(⋅|c)x_T p_0(·|c) (resp. pu⋆(⋅|c)p_u (·|c)). Then when Df=DKLD_f=D_ KL, we have qu⋆,t(xt|c) q_u ,t(x_t|c) =∫qt(xt|xT)pu⋆(xT|c)xT = q_t(x_t|x_T)p_u (x_T|c)dx_T =(61)1Z0(c)∫qt(xt|xT)p0(xT|c)exp(1τr(xT,c))xT eq:p_u_star_KL= 1Z_0(c) q_t(x_t|x_T)p_0(x_T|c) ( 1τr (x_T,c ) )dx_T =1Z0(c)∫q0,t(xT|xt,c)q0,t(xt|c)exp(1τr(xT,c))xT = 1Z_0(c) q_0,t(x_T|x_t,c)q_0,t(x_t|c) ( 1τr (x_T,c ) )dx_T =q0,t(xt|c)xT∼q0,t(⋅|xt,c)[exp(1τr(xT,c))]Z0(c), =q_0,t(x_t|c) E_x_T q_0,t(·|x_t,c) [ ( 1τr (x_T,c ) ) ]Z_0(c), (63) where the third line follows from the following relation given by the Bayes’ rule: q0,t(xT|xt,c)=qt(xt|xT)p0(xT|c)q0,t(xt|c). q_0,t(x_T|x_t,c)= q_t(x_t|x_T)p_0(x_T|c)q_0,t(x_t|c). (64) Thus we have qu⋆,t(xt+1|xt,c) q_u ,t(x_t+1|x_t,c) =qt+1(xt|xt+1)qu⋆,t(xt+1|c)qu⋆,t(xt|c) = q_t+1(x_t|x_t+1)q_u ,t(x_t+1|c)q_u ,t(x_t|c) =(B.3)qt+1(xt|xt+1)q0,t+1(xt+1|c)q0,t(xt|c)xT∼q0,t+1(⋅|xt+1,c)[exp(1τr(xT,c))]xT∼q0,t(⋅|xt,c)[exp(1τr(xT,c))] eq:q_u_star_t= q_t+1(x_t|x_t+1)q_0,t+1(x_t+1|c)q_0,t(x_t|c) E_x_T q_0,t+1(·|x_t+1,c) [ ( 1τr (x_T,c ) ) ]E_x_T q_0,t(·|x_t,c) [ ( 1τr (x_T,c ) ) ] =q0,t(xt+1|xt,c)xT∼q0,t+1(⋅|xt+1,c)[exp(1τr(xT,c))]xT∼q0,t(⋅|xt,c)[exp(1τr(xT,c))], =q_0,t(x_t+1|x_t,c) E_x_T q_0,t+1(·|x_t+1,c) [ ( 1τr (x_T,c ) ) ]E_x_T q_0,t(·|x_t,c) [ ( 1τr (x_T,c ) ) ], (65) where q0,t(xt+1|xt,c)q_0,t(x_t+1|x_t,c) and qu⋆,t(xt+1|xt,c)q_u ,t(x_t+1|x_t,c) represent the posterior conditional distributions induced by the forward process starting from target distributions p0p_0 and pu⋆p_u , respectively. We define ψt(xt,c)≔xT∼q0,t(⋅|xt,c)[exp(1τr(xT,c))], _t(x_t,c) _x_T q_0,t(·|x_t,c) [ ( 1τr (x_T,c ) ) ], (66) then by tower property of expectation, we have ψt(xt,c)=xt+1∼q0,t(⋅|xt,c)[ψt+1(xt+1,c)], _t(x_t,c)=E_x_t+1 q_0,t(·|x_t,c) [ _t+1(x_t+1,c) ], (67) and by (B.3), we have qu⋆,t(xt+1|xt,c)=q0,t(xt+1|xt,c)ψt+1(xt+1,c)ψt(xt,c). q_u ,t(x_t+1|x_t,c)=q_0,t(x_t+1|x_t,c) _t+1(x_t+1,c) _t(x_t,c). (68) Thus we can express the mean of qu⋆,t(xt+1|xt,c)q_u ,t(x_t+1|x_t,c) as follows: xt+1∼qu⋆,t(⋅|xt,c)[xt+1]=xt+1∼q0,t(⋅|xt,c)[xt+1ψt+1(xt+1,c)]ψt(xt,c). _x_t+1 q_u ,t(·|x_t,c) [x_t+1 ]= E_x_t+1 q_0,t(·|x_t,c) [x_t+1 _t+1(x_t+1,c) ] _t(x_t,c). (69) We let pu⋆,t(xt+1|xt,c)p_u ,t(x_t+1|x_t,c) denote the reverse conditional distribution of xt+1x_t+1 given xtx_t under the optimal score function ϵ⋆ε . Then we have the following lemma that says the mean of the posterior conditional distribution qu⋆,t(⋅|xt,c)q_u ,t(·|x_t,c) induced by the forward process is the same as the mean of the learned reverse kernel pu⋆,t(⋅|xt,c)p_u ,t(·|x_t,c): Lemma 2 For any data distribution p(⋅|c)p(·|c) of xTx_T (∀c∈∀ c ), let ϵp⋆ _p be the optimal score function that minimizes the diffusion loss ℒorg(⋅;p)L_org (·;p ) defined in (10). Then for any xt∈ℝd,c∈,t∈[T−1]x_t ^d,c ,t∈[T-1], we have xt+1∼qt(⋅|xt,c)[xt+1]=μp⋆(xt,c,t)≔1αtxt−βtαt(1−α¯t)ϵp⋆(xt,c,t), _x_t+1 q_t(·|x_t,c) [x_t+1 ]= _p (x_t,c,t) 1 _tx_t- _t _t (1- α_t ) _p (x_t,c,t), (70) where qt(⋅|xt,c)q_t(·|x_t,c) is the posterior conditional distribution induced by the forward process starting from data distribution p. The proof of Lemma 2 is given in Appendix B.5. By Lemma 2, we have μ⋆(xt,c,t) μ (x_t,c,t) =xt+1∼qu⋆,t(⋅|xt,c)[xt+1] =E_x_t+1 q_u ,t(·|x_t,c) [x_t+1 ] =(69)xt+1∼q0,t(⋅|xt,c)[xt+1ψt+1(xt+1,c)]ψt(xt,c) eq:mu_u_star_t= E_x_t+1 q_0,t(·|x_t,c) [x_t+1 _t+1(x_t+1,c) ] _t(x_t,c) =xt+1∼p0,t(⋅|xt,c)[xt+1ψt+1(xt+1,c)]xt+1∼p0,t(⋅|xt,c)[ψt+1(xt+1,c)]⏟≔μ~⋆(xt,c,t) = E_x_t+1 p_0,t(·|x_t,c) [x_t+1 _t+1(x_t+1,c) ]E_x_t+1 p_0,t(·|x_t,c) [ _t+1(x_t+1,c) ]_ μ (x_t,c,t) +xt+1∼q0,t(⋅|xt,c)[xt+1ψt+1(xt+1,c)]ψt(xt,c)−xt+1∼p0,t(⋅|xt,c)[xt+1ψt+1(xt+1,c)]xt+1∼p0,t(⋅|xt,c)[ψt+1(xt+1,c)]⏟≔δ0(xt,c,t). + E_x_t+1 q_0,t(·|x_t,c) [x_t+1 _t+1(x_t+1,c) ] _t(x_t,c)- E_x_t+1 p_0,t(·|x_t,c) [x_t+1 _t+1(x_t+1,c) ]E_x_t+1 p_0,t(·|x_t,c) [ _t+1(x_t+1,c) ]_ _0(x_t,c,t). (71) Under this decomposition and the relation between μ and ϵε, we have ϵ⋆(xt,c,t)−ϵ0(xt,c,t) ε (x_t,c,t)- _0(x_t,c,t) =αt(1−α¯t)βt(μ0(xt,c,t)−μ⋆(xt,c,t)) = _t (1- α_t ) _t ( _0(x_t,c,t)-μ (x_t,c,t) ) =αt(1−α¯t)βt(μ0(xt,c,t)−μ~⋆(xt,c,t)−δ0(xt,c,t)). = _t (1- α_t ) _t ( _0(x_t,c,t)- μ (x_t,c,t)- _0(x_t,c,t) ). (72) Define ψ~t(xt,c)≔xt+1∼p0,t(⋅|xt,c)[ψt+1(xt+1,c)], ψ_t(x_t,c) _x_t+1 p_0,t(·|x_t,c) [ _t+1(x_t+1,c) ], (73) then combining (B.3) and (B.3), we have ϵ⋆(xt,c,t)−ϵ0(xt,c,t) ε (x_t,c,t)- _0(x_t,c,t) =αt(1−α¯t)βt(−xt+1∼p0,t(⋅|xt,c)[xt+1(ψt+1(xt+1,c)−1)]ψ~t(xt,c)+(1−1ψ~t(xt,c))μ0(xt,c,t)−δ0(xt,c,t)), = _t (1- α_t ) _t (- E_x_t+1 p_0,t(·|x_t,c) [x_t+1( _t+1(x_t+1,c)-1) ] ψ_t(x_t,c)+ (1- 1 ψ_t(x_t,c) ) _0(x_t,c,t)- _0(x_t,c,t) ), (74) where we use the fact that μ0(xt,c,t)=xt+1∼p0,t(⋅|xt,c)[xt+1]. _0(x_t,c,t)=E_x_t+1 p_0,t(·|x_t,c) [x_t+1 ]. (75) Recall in (23) we define the basis functions ϕω(xt,c,t) _ω(x_t,c,t) as ϕω(xt,c,t)=(cos(ω⊤μ0(xt,c,t)β~t),sin(ω⊤μ0(xt,c,t)β~t))⊤, _ω(x_t,c,t)= ( ( ω _0(x_t,c,t) β_t ), ( ω _0(x_t,c,t) β_t ) ) , (76) and also define ρω(xt+1,t)=1(2πβ~t)d(cos(ω⊤xt+1β~t),sin(ω⊤xt+1β~t))⊤. _ω(x_t+1,t)= 1 (2π β_t)^d ( ( ω x_t+1 β_t ), ( ω x_t+1 β_t ) ) . (77) Then we can decompose the Gaussian reverse kernel p0,t(xt+1|xt,c)p_0,t (x_t+1|x_t,c ) as p0,t(xt+1|xt,c)=⟨ϕω(xt,t),ρω(xt+1,t)⟩(ω), p_0,t (x_t+1|x_t,c )= _ω(x_t,t), _ω(x_t+1,t) _N (ω ), (78) where ⟨fω,gω⟩(ω)≔ω∼(⋅|0,d)[fω⊤gω] f_ω,g_ω _N (ω ) _ω (·|0,I_d ) [f_ω g_ω ] for any fωf_ω and gωg_ω. We define h(xt,c,t) h(x_t,c,t) ≔xt+1∼p0,t(⋅|xt,c)[xt+1(ψt+1(xt+1,c)−1)], _x_t+1 p_0,t(·|x_t,c) [x_t+1( _t+1(x_t+1,c)-1) ], (79) z(xt,c,t) z(x_t,c,t) ≔1−1ψ~t(xt,c), 1- 1 ψ_t(x_t,c), (80) then by (B.3) and (7), ϵ⋆(xt,c,t) ε (x_t,c,t) =1ψ~t(xt,c)ϵ0(xt,c,t)−1−α¯tβt(αtψ~t(xt,c)h(xt,c,t)+(1−1ψ~t(xt,c))xt−αtδ0(xt,c,t)) = 1 ψ_t(x_t,c) _0(x_t,c,t)- 1- α_t _t ( _t ψ_t(x_t,c)h(x_t,c,t)+ (1- 1 ψ_t(x_t,c) )x_t- _t _0(x_t,c,t) ) =(1−z(xt,c,t))ϵ0(xt,c,t)−1−α¯tβt(αt(1−z(xt,c,t))h(xt,c,t)+z(xt,c,t)xt−αtδ0(xt,c,t)). =(1-z(x_t,c,t)) _0(x_t,c,t)- 1- α_t _t ( _t(1-z(x_t,c,t))h(x_t,c,t)+z(x_t,c,t)x_t- _t _0(x_t,c,t) ). (81) Moreover, we can express h(xt,c,t)h(x_t,c,t) as h(xt,c,t)=ω∼(0,d)[(∫(ψt+1(xt+1,c)−1)xt+1ρω(xt+1,t)⊤xt+1⏟W⋆(ω,c,t)∈ℝd×2)ϕω(xt,c,t)], h(x_t,c,t)=E_ω (0,I_d) [ ( ( _t+1(x_t+1,c)-1 )x_t+1 _ω(x_t+1,t) dx_t+1_W (ω,c,t) ^d× 2 ) _ω(x_t,c,t) ], (82) and express ψ~t(xt,c)−1=z(xt,c,t)1−z(xt,c,t) ψ_t(x_t,c)-1= z(x_t,c,t)1-z(x_t,c,t) as z(xt,c,t)1−z(xt,c,t)=ψ~t(xt,c)−1=(73)ω∼(0,d)[(∫(ψt+1(xt+1,c)−1)ρω(xt+1,t)⊤xt+1⏟u⋆(ω,c,t)∈ℝ2)ϕω(xt,c,t)], z(x_t,c,t)1-z(x_t,c,t)= ψ_t(x_t,c)-1 eq:widetilde_psi_t=E_ω (0,I_d) [ ( ( _t+1(x_t+1,c)-1 ) _ω(x_t+1,t) dx_t+1_u (ω,c,t) ^2 ) _ω(x_t,c,t) ], (83) We can use random Fourier features to approximate them as h(xt,c,t) h(x_t,c,t) =1M∑i=1MW⋆(ωi,c,t)ϕωi(xt,c,t)+δh(xt,c,t), = 1M _i=1^MW ( _i,c,t) _ _i(x_t,c,t)+ _h(x_t,c,t), (84) z(xt,c,t)1−z(xt,c,t) z(x_t,c,t)1-z(x_t,c,t) =1M∑i=1Mu⋆(ωi,c,t)⊤ϕωi(xt,c,t)+δz(xt,c,t), = 1M _i=1^Mu ( _i,c,t) _ _i(x_t,c,t)+ _z(x_t,c,t), (85) where ωi∼i.i.d.(0,d) _i i.i.d. N(0,I_d), i∈[M]i∈[M]. Let BW B_W ≔esssupω∈ℝd‖W⋆(ω,c,t)‖F,VW2≔esssupω∈ℝd‖u⋆(ω,c,t)‖F2, _ω ^d \|W (ω,c,t) \|_F, V_W^2 _ω ^d \|u (ω,c,t) \|_F^2, Bu B_u ≔esssupω∈ℝd‖u⋆(ω,c,t)‖2,Vu2≔esssupω∈ℝd‖u⋆(ω,c,t)‖22. _ω ^d \|u (ω,c,t) \|_2, V_u^2 _ω ^d \|u (ω,c,t) \|_2^2. (86) By Bernstein’s inequality (Kwon and Perchet [2017, Proposition E.3]), we have with probability at least 1−δ1-δ, ‖δh‖2⩽2VW2log(2/δ)M+2BWlog(2/δ)3M,‖δz‖2⩽2Vu2log(2/δ)M+2Bulog(2/δ)3M. \| _h \|_2 2V_W^2 (2/δ)M+ 2B_W (2/δ)3M, \| _z \|_2 2V_u^2 (2/δ)M+ 2B_u (2/δ)3M. (87) Finally, we bound δ0(xt,c,t) _0(x_t,c,t) defined in (B.3). Recall we assume xT∈∞d(1)x_T _∞^d(1), where we let ∞d(1)≔xT∈ℝd:‖xT‖∞⩽1. _∞^d(1) \x_T ^d: \|x_T \|_∞ 1 \. (88) We let rmin≔infxT∈∞d(1),c∈r(xT,c),rmax≔supxT∈∞d(1),c∈r(xT,c). r_min _x_T _∞^d(1),c r(x_T,c), r_max _x_T _∞^d(1),c r(x_T,c). (89) Then rminr_min and rmaxr_max are both finite, and by (66) we know that ψmin≔exp(rminτ)⩽ψt(xt,c)⩽ψmax≔exp(rmaxτ). _min ( r_minτ ) _t(x_t,c) _max ( r_maxτ ). (90) we define weighting function wt(xt,c)≔ψt(xt,c)ψmin,w¯t(xt,c)≔wt(xt,c)−1, w_t(x_t,c) _t(x_t,c) _min, w_t(x_t,c) w_t(x_t,c)-1, (91) Then we have δ0(xt,c,t) _0(x_t,c,t) ≔xt+1∼q0,t(⋅|xt,c)[xt+1wt+1(xt+1,c)]xt+1∼q0,t(⋅|xt,c)[wt+1(xt+1,c)]−xt+1∼p0,t(⋅|xt,c)[xt+1wt+1(xt+1,c)]xt+1∼p0,t(⋅|xt,c)[wt+1(xt+1,c)] E_x_t+1 q_0,t(·|x_t,c) [x_t+1w_t+1(x_t+1,c) ]E_x_t+1 q_0,t(·|x_t,c) [w_t+1(x_t+1,c) ]- E_x_t+1 p_0,t(·|x_t,c) [x_t+1w_t+1(x_t+1,c) ]E_x_t+1 p_0,t(·|x_t,c) [w_t+1(x_t+1,c) ] =xt+1∼q0,t(⋅|xt,c)[(xt+1−μ0(xt,c,t))wt+1(xt+1,c)]xt+1∼q0,t(⋅|xt,c)[wt+1(xt+1,c)]−xt+1∼p0,t(⋅|xt,c)[(xt+1−μ0(xt,c,t))wt+1(xt+1,c)]xt+1∼p0,t(⋅|xt,c)[wt+1(xt+1,c)] = E_x_t+1 q_0,t(·|x_t,c) [(x_t+1- _0(x_t,c,t))w_t+1(x_t+1,c) ]E_x_t+1 q_0,t(·|x_t,c) [w_t+1(x_t+1,c) ]- E_x_t+1 p_0,t(·|x_t,c) [(x_t+1- _0(x_t,c,t))w_t+1(x_t+1,c) ]E_x_t+1 p_0,t(·|x_t,c) [w_t+1(x_t+1,c) ] =xt+1∼q0,t(⋅|xt,c)[(xt+1−μ0(xt,c,t))w¯t+1(xt+1,c)]xt+1∼q0,t(⋅|xt,c)[wt+1(xt+1,c)]−xt+1∼p0,t(⋅|xt,c)[(xt+1−μ0(xt,c,t))w¯t+1(xt+1,c)]xt+1∼p0,t(⋅|xt,c)[wt+1(xt+1,c)] = E_x_t+1 q_0,t(·|x_t,c) [(x_t+1- _0(x_t,c,t)) w_t+1(x_t+1,c) ]E_x_t+1 q_0,t(·|x_t,c) [w_t+1(x_t+1,c) ]- E_x_t+1 p_0,t(·|x_t,c) [(x_t+1- _0(x_t,c,t)) w_t+1(x_t+1,c) ]E_x_t+1 p_0,t(·|x_t,c) [w_t+1(x_t+1,c) ] (92) where the first line uses (B.3) and (67), and the third line follows from the fact that μ0(xt,c,t)=xt+1∼p0,t(⋅|xt,c)[xt+1]=xt+1∼q0,t(⋅|xt,c)[xt+1] _0(x_t,c,t)=E_x_t+1 p_0,t(·|x_t,c) [x_t+1 ]=E_x_t+1 q_0,t(·|x_t,c) [x_t+1 ] (93) given by the optimality of ϵ0 _0 and Lemma 2. Define Aq,t≔xt+1∼q0,t(⋅|xt,c)[(xt+1−μ0(xt,c,t))w¯t+1(xt+1,c)],Ap,t≔xt+1∼p0,t(⋅|xt,c)[(xt+1−μ0(xt,c,t))w¯t+1(xt+1,c)],Bq,t≔xt+1∼q0,t(⋅|xt,c)[wt+1(xt+1,c)],Bp,t≔xt+1∼p0,t(⋅|xt,c)[wt+1(xt+1,c)], splitA_q,t& _x_t+1 q_0,t(·|x_t,c) [(x_t+1- _0(x_t,c,t)) w_t+1(x_t+1,c) ],\\ A_p,t& _x_t+1 p_0,t(·|x_t,c) [(x_t+1- _0(x_t,c,t)) w_t+1(x_t+1,c) ],\\ B_q,t& _x_t+1 q_0,t(·|x_t,c) [w_t+1(x_t+1,c) ],\\ B_p,t& _x_t+1 p_0,t(·|x_t,c) [w_t+1(x_t+1,c) ], split (94) Then ‖δ0(xt,c,t)‖2 \| _0(x_t,c,t) \|_2 =‖Aq,tBq,t−Ap,tBp,t‖2 = \| A_q,tB_q,t- A_p,tB_p,t \|_2 ⩽‖Aq,t−Ap,tBq,t‖2+‖Ap,t‖2|1Bp,t−1Bq,t| \| A_q,t-A_p,tB_q,t \|_2+ \|A_p,t \|_2 | 1B_p,t- 1B_q,t | ⩽‖Aq,t−Ap,t‖2⏟(i)+‖Ap,t‖2|Bp,t−Bq,t|⏟(ii), \|A_q,t-A_p,t \|_2_(i)+ \|A_p,t \|_2 |B_p,t-B_q,t |_(i), (95) where we use the fact that Bp,t⩾1,Bq,t⩾1B_p,t 1, B_q,t 1 by our choice of wt+1w_t+1 (c.f. (91)). Below we bound (i)(i) and (ii)(i) in (B.3) separately utilizing the standard TV inequality: ‖p[g]−q[g]‖2⩽2TV(p,q)p[‖g‖22]+q[‖g‖22]. \|E_p [g ]-E_q [g ] \|_2 2TV (p,q ) E_p [ \|g \|_2^2 ]+E_q [ \|g \|_2^2 ]. (96) for any vector-valued function g and probability measure p and q. Define gt+1≔(xt+1−μ0(xt,c,t))w¯t+1(xt+1,c), g_t+1 (x_t+1- _0(x_t,c,t)) w_t+1(x_t+1,c), (97) then by (91) we have ‖gt+1‖22⩽(ψmaxψmin−1)2‖xt+1−μ0(xt,c,t)‖22. \|g_t+1 \|_2^2 ( _max _min-1 )^2 \|x_t+1- _0(x_t,c,t) \|_2^2. (98) and by (96) we have (i) (i) =‖xt+1∼q0,t(⋅|xt,c)[gt+1]−xt+1∼p0,t(⋅|xt,c)[gt+1]‖2 = \|E_x_t+1 q_0,t(·|x_t,c) [g_t+1 ]-E_x_t+1 p_0,t(·|x_t,c) [g_t+1 ] \|_2 ⩽2TV(q0,t,p0,t)xt+1∼q0,t(⋅|xt,c)[‖gt+1‖22]+xt+1∼p0,t(⋅|xt,c)[‖gt+1‖22] 2TV (q_0,t,p_0,t ) E_x_t+1 q_0,t(·|x_t,c) [ \|g_t+1 \|_2^2 ]+E_x_t+1 p_0,t(·|x_t,c) [ \|g_t+1 \|_2^2 ] ⩽(98)(ψmaxψmin−1)2TV(q0,t,p0,t) eq:||g_t+1||^2_2 ( _max _min-1 ) 2TV (q_0,t,p_0,t ) ⋅xt+1∼q0,t(⋅|xt,c)[‖xt+1−μ0(xt,c,t)‖22]+xt+1∼p0,t(⋅|xt,c)[‖xt+1−μ0(xt,c,t)‖22], · E_x_t+1 q_0,t(·|x_t,c) [ \|x_t+1- _0(x_t,c,t) \|_2^2 ]+E_x_t+1 p_0,t(·|x_t,c) [ \|x_t+1- _0(x_t,c,t) \|_2^2 ], (99) where we write TV(q0,t,p0,t)TV (q_0,t,p_0,t ) as a shorthand for TV(q0,t(⋅|xt,c),p0,t(⋅|xt,c))TV (q_0,t(·|x_t,c),p_0,t(·|x_t,c) ). For (i), by (98) we have ‖Ap,t‖2⩽(ψmaxψmin−1)xt+1∼p0,t(⋅|xt,c)[‖xt+1−μ0(xt,c,t)‖22], \|A_p,t \|_2 ( _max _min-1 ) E_x_t+1 p_0,t(·|x_t,c) [ \|x_t+1- _0(x_t,c,t) \|_2^2 ], (100) and by (96) we have |Bp,t−Bq,t| |B_p,t-B_q,t| ⩽2TV(q0,t,p0,t)xt+1∼q0,t(⋅|xt,c)[‖wt+1(xt+1,c)‖22]+xt+1∼p0,t(⋅|xt,c)[‖wt+1(xt+1,c)‖22] 2TV (q_0,t,p_0,t ) E_x_t+1 q_0,t(·|x_t,c) [ \|w_t+1(x_t+1,c) \|_2^2 ]+E_x_t+1 p_0,t(·|x_t,c) [ \|w_t+1(x_t+1,c) \|_2^2 ] ⩽2ψmaxψmin2TV(q0,t,p0,t). 2 _max _min 2TV (q_0,t,p_0,t ). (101) Combining The above two expressions, we have (ii) (i) ⩽2ψmaxψmin(ψmaxψmin−1)xt+1∼q0,t(⋅|xt,c)[‖xt+1−μ0(xt,c,t)‖22]2TV(q0,t,p0,t) 2 _max _min ( _max _min-1 ) E_x_t+1 q_0,t(·|x_t,c) [ \|x_t+1- _0(x_t,c,t) \|_2^2 ] 2TV (q_0,t,p_0,t ) (102) Combining (B.3) and (102), we have ‖δ0(xt,c,t)‖2 \| _0(x_t,c,t) \|_2 ⩽(ψmaxψmin−1)(1+2ψmaxψmin)2TV(q0,t,p0,t) ( _max _min-1 ) (1+ 2 _max _min ) 2TV (q_0,t,p_0,t ) ⋅xt+1∼q0,t(⋅|xt,c)[‖xt+1−μ0(xt,c,t)‖22]⏟(a)+xt+1∼p0,t(⋅|xt,c)[‖xt+1−μ0(xt,c,t)‖22]⏟(b). · E_x_t+1 q_0,t(·|x_t,c) [ \|x_t+1- _0(x_t,c,t) \|_2^2 ]_(a)+ E_x_t+1 p_0,t(·|x_t,c) [ \|x_t+1- _0(x_t,c,t) \|_2^2 ]_(b). (103) For (b), since q0,t(⋅|xt,c)=(⋅|μ0(xt,c,t),β~td)q_0,t(·|x_t,c)=N (· | _0(x_t,c,t), β_tI_d ), we have (b) (b) ≔xt+1∼p0,t(⋅|xt,c)[‖xt+1−μ0(xt,c,t)‖22]=dβt~, _x_t+1 p_0,t(·|x_t,c) [ \|x_t+1- _0(x_t,c,t) \|_2^2 ]=d _t, (104) For (a), we have (a) (a) =xt+1∼q0,t(⋅|xt,c)[‖xt+1−μ0(xt,c,t)‖22] =E_x_t+1 q_0,t(·|x_t,c) [ \|x_t+1- _0(x_t,c,t) \|_2^2 ] =xT∼q0,t(⋅|xt,c)[xt+1∼q0,t(⋅|xt,xT,c)[‖xt+1−μ0(xt,c,t)‖22]]. =E_x_T q_0,t(·|x_t,c) [E_x_t+1 q_0,t(·|x_t,x_T,c) [ \|x_t+1- _0(x_t,c,t) \|_2^2 ] ]. (105) By (114) we can compute the inner expectation as xt+1∼q0,t(⋅|xt,xT,c)[‖xt+1−μ0(xt,c,t)‖22] _x_t+1 q_0,t(·|x_t,x_T,c) [ \|x_t+1- _0(x_t,c,t) \|_2^2 ] =ξ∼(0,d)[‖μ~t(xt,xT,t)−μ0(xt,c,t)+β~tξ‖22] =E_ξ (0,I_d ) [ \| μ_t(x_t,x_T,t)- _0(x_t,c,t)+ β_tξ \|_2^2 ] =‖μ~t(xt,xT,t)−μ0(xt,c,t)‖22+dβ~t, = \| μ_t(x_t,x_T,t)- _0(x_t,c,t) \|_2^2+d β_t, (106) where β~t=βt(1−α¯t+1)1−α¯t,μ~t(xt,xT,t)=αt(1−α¯t+1)1−α¯txt+βtα¯t+11−α¯txT. β_t= _t(1- α_t+1)1- α_t, μ_t(x_t,x_T,t)= _t(1- α_t+1)1- α_tx_t+ _t α_t+11- α_tx_T. (107) Plugging (B.3) back into (B.3), we have (a) (a) =xT∼q0,t(⋅|xt,c)[‖μ~t(xt,xT,t)−μ0(xt,c,t)‖22]+dβ~t =E_x_T q_0,t(·|x_t,c) [ \| μ_t(x_t,x_T,t)- _0(x_t,c,t) \|_2^2 ]+d β_t =xT∼q0,t(⋅|xt,c)[‖μ~t(xt,xT,t)−xT∼q0,t(⋅|xt,c)[μ~t(xt,xT,t)]‖22]+dβ~t =E_x_T q_0,t(·|x_t,c) [ \| μ_t(x_t,x_T,t)-E_x_T q_0,t(·|x_t,c) [ μ_t(x_t,x_T,t) ] \|_2^2 ]+d β_t =(107)(βtα¯t+11−α¯t)2xT∼qt(⋅|xt,c)[∥xT−[xT|xt,c]∥22]+dβ~t, eq:tilde_beta_t_mu_t= ( _t α_t+11- α_t )^2E_x_T q_t(·|x_t,c) [ \|x_T-E [x_T|x_t,c ] \|_2^2 ]+d β_t, (108) where the second equality follows from Lemma 2 and (117) in its proof. Plugging (B.3) and (104) back into (B.3), we have ‖δ0(xt,c,t)‖2 \| _0(x_t,c,t) \|_2 ⩽(ψmaxψmin−1)(1+2ψmaxψmin)2TV(q0,t,p0,t) ( _max _min-1 ) (1+ 2 _max _min ) 2TV (q_0,t,p_0,t ) ⋅(βtα¯t+11−α¯t)2xT∼q0,t(⋅|xt,c)[∥xT−[xT|xt,c]∥22]+d(β~t+βt) · ( _t α_t+11- α_t )^2E_x_T q_0,t(·|x_t,c) [ \|x_T-E [x_T|x_t,c ] \|_2^2 ]+d( β_t+ _t) ⩽(ψmaxψmin−1)(1+2ψmaxψmin)2TV(q0,t,p0,t)(βtα¯t+11−α¯t)2+d(β~t+βt), ( _max _min-1 ) (1+ 2 _max _min ) 2TV (q_0,t,p_0,t ) ( _t α_t+11- α_t )^2+d( β_t+ _t), (109) where the last inequality follows from our assumption that ‖xT‖∞⩽1 \|x_T \|_∞ 1. B.4 Proof of Lemma 1 Note that ∥ϵ(xt,c,t)+1−α¯t∇xtlogqt⋆(xt|c)∥22 \|ε (x_t,c,t )+ 1- α_t _x_t q _t (x_t|c ) \|_2^2 =∥ϵ(xt,c,t)+1−α¯tqt⋆(xt|c)∇xt∫p⋆(xT|c)qt(xt|xT)dxT∥22 = \|ε (x_t,c,t )+ 1- α_tq _t (x_t|c ) _x_t p (x_T|c )q_t(x_t|x_T)dx_T \|_2^2 =‖ϵ(xt,c,t)‖22+21−α¯tqt⋆(xt|c)⟨ϵ(xt,c,t),∇xt∫p⋆(xT|c)qt(xt|xT)xT⟩+ = \|ε (x_t,c,t ) \|_2^2+ 2 1- α_tq _t (x_t|c ) ε (x_t,c,t ), _x_t p (x_T|c )q_t(x_t|x_T)dx_T + const =‖ϵ(xt,c,t)‖221qt⋆(xt|c)∫p⋆(xT|c)qt(xt|xT)xT⏟=1 = \|ε (x_t,c,t ) \|_2^2 1q _t (x_t|c ) p (x_T|c )q_t(x_t|x_T)dx_T_=1 +21−α¯tqt⋆(xt|c)∫p⋆(xT|c)qt(xt|xT)⟨ϵ(xt,c,t),∇xtlogqt(xt|xT)⟩xT+ + 2 1- α_tq _t (x_t|c ) p (x_T|c )q_t(x_t|x_T) ε (x_t,c,t ), _x_t q_t(x_t|x_T) dx_T+ const =1qt⋆(xt|c)∫p⋆(xT|c)qt(xt|xT)(‖ϵ(xt,c,t)‖22+21−α¯t⟨ϵ(xt,c,t),∇xtlogqt(xt|xT)⟩)xT+ = 1q _t (x_t|c ) p (x_T|c )q_t(x_t|x_T) ( \|ε (x_t,c,t ) \|_2^2+2 1- α_t ε (x_t,c,t ), _x_t q_t(x_t|x_T) )dx_T+ const =1qt⋆(xt|c)∫p⋆(xT|c)qt(xt|xT)(∥ϵ(xt,c,t)+1−α¯t∇xtlogqt(xt|xT)∥22)dxT+, = 1q _t (x_t|c ) p (x_T|c )q_t(x_t|x_T) ( \|ε (x_t,c,t )+ 1- α_t _x_t q_t(x_t|x_T) \|_2^2 )dx_T+ const, (110) where const denotes constants that are independent of ϵε. Integrate both sides of the above equation over xt∼qt⋆(⋅|c)x_t q _t (·|c ), we have xt∼qt⋆(⋅|c)[∥ϵ(xt,c,t)+1−α¯t∇xtlogqt⋆(xt|c)∥22] _x_t q _t (·|c ) [ \|ε (x_t,c,t )+ 1- α_t _x_t q _t (x_t|c ) \|_2^2 ] =xT∼p⋆(⋅|c),xt∼qt(⋅|xT)[∥ϵ(xt,c,t)+1−α¯t∇xtlogqt(xt|xT)∥22]+ =E_ subarraycx_T p (·|c ),\\ x_t q_t(·|x_T) subarray [ \|ε (x_t,c,t )+ 1- α_t _x_t q_t(x_t|x_T) \|_2^2 ]+ const =xT∼p⋆(⋅|c),ξ∼(0,d)[‖ξ−ϵ(α¯txT+1−α¯tξ,c,t)‖22]+. =E_ subarraycx_T p (·|c ),\\ ξ (0,I_d ) subarray [ \|ξ-ε ( α_tx_T+ 1- α_tξ,c,t ) \|_2^2 ]+ const. (111) By comparing (B.4) with ℒorg(⋅;p⋆)L_ org(·;p ), we obtain ℒorg(ϵ;p⋆)=12c∼pc,t∼([T−1]),xt∼qt⋆(⋅|c)[∥ϵ(xt,c,t)+1−α¯t∇xtlogqt⋆(xt|c)∥22]+. _ org (ε;p )= 12E_ subarraycc p_c,t ([T-1] ),\\ x_t q _t (·|c ) subarray [ \|ε (x_t,c,t )+ 1- α_t _x_t q _t (x_t|c ) \|_2^2 ]+ const. (112) Recall we assume supp(pc)= supp (p_c )=C. Therefore, (112) indicates the minimum of ℒorg(⋅;p⋆)L_ org(·;p ) is ϵp⋆ _p that satisfies (42). On the other hand, it’s obvious that for any g:ℝd×[T]→(0,+∞)g:R^d×C×[T]→(0,+∞), the minimum of ℒL is also ϵp⋆ _p , as long as (41) is well-defined. B.5 Proof of Lemma 2 We let μq(xt,c,t)≔xt+1∼qt(⋅|xt,c)[xt+1]. _q(x_t,c,t) _x_t+1 q_t(·|x_t,c) [x_t+1 ]. (113) First note that by the forward process rule (4) and Bayes’ rule, we have qt(xt+1|xt,xT)=qt+1(xt+1|xT)pt(xt|xt+1)qt(xt|xT)=(xt+1|μ~t(xt,xT,t),β~td) q_t(x_t+1|x_t,x_T)= q_t+1(x_t+1|x_T)p_t(x_t|x_t+1)q_t(x_t|x_T)=N (x_t+1| μ_t(x_t,x_T,t), β_tI_d ) (114) with μ~t(xt,xT,t)=αt(1−α¯t+1)1−α¯txt+βtα¯t+11−α¯txT μ_t(x_t,x_T,t)= _t(1- α_t+1)1- α_tx_t+ _t α_t+11- α_tx_T (115) and β~t=βt(1−α¯t+1)1−α¯t. β_t= _t(1- α_t+1)1- α_t. (116) Thus we have μq(xt,c,t)=xT∼qt(⋅|xt,c)[μ~t(xt,xT,t)]=αt(1−α¯t+1)1−α¯txt+βtα¯t+11−α¯txT∼qt(⋅|xt,c)[xT]. _q(x_t,c,t)=E_x_T q_t(·|x_t,c) [ μ_t(x_t,x_T,t) ]= _t(1- α_t+1)1- α_tx_t+ _t α_t+11- α_tE_x_T q_t(·|x_t,c) [x_T ]. (117) On the other hand, since ϵp⋆ _p minimizes the noise-prediction regression objective ℒ0(⋅;p)L_0 (·;p ) (c.f. (10)), we have ϵp⋆(xt,c,t)=[ξ|xt,c,t]=(4)[xt−α¯txT1−α¯t|xt,c,t]=xt1−α¯t−α¯t1−α¯txT∼qt(⋅|xt,c)[xT], _p (x_t,c,t)=E[ξ|x_t,c,t] eq:forward_process_1=E [ x_t- α_tx_T 1- α_t |x_t,c,t ]= x_t 1- α_t- α_t 1- α_tE_x_T q_t(·|x_t,c) [x_T ], (118) from which we deduce μp⋆(xt,c,t) _p (x_t,c,t) =1αtxt−βtαt(1−α¯t)ϵp⋆(xt,c,t) = 1 _tx_t- _t _t (1- α_t ) _p (x_t,c,t) =1αtxt−βtαt(1−α¯t)(xt1−α¯t−α¯t1−α¯txT∼qt(⋅|xt,c)[xT]) = 1 _tx_t- _t _t (1- α_t ) ( x_t 1- α_t- α_t 1- α_tE_x_T q_t(·|x_t,c) [x_T ] ) =αt(1−α¯t+1)1−α¯txt+βtα¯t+11−α¯txT∼qt(⋅|xt,c)[xT]=(117)μq(xt,c,t). = _t(1- α_t+1)1- α_tx_t+ _t α_t+11- α_tE_x_T q_t(·|x_t,c) [x_T ] eq:mu_q= _q(x_t,c,t). (119) Appendix C Additional Experiment Details Details on evaluation. For automatic evaluation, we found at test time, enlarging the side network guidance strength λmodel _ model can help the finetuned models to achieve better performance, especially for SFT and RWL. Thus for SFT and RWL, we evaluate DiffCon, DiffCon-J, DiffCon-S with 1/λmodel∈0.1,0.2,⋯,1.01/ _ model∈\0.1,0.2,·s,1.0\ and DiffCon-Naive with λmodel∈1.0,2.0,⋯,10.0 _ model∈\1.0,2.0,·s,10.0\; for PPO, we evaluate DiffCon, DiffCon-J, DiffCon-S with 1/λmodel∈0.7,0.8,0.9,1.01/ _ model∈\0.7,0.8,0.9,1.0\ and DiffCon-Naive with λmodel∈1.0,2.0,3.0,4.0 _ model∈\1.0,2.0,3.0,4.0\, and we report the best HPS-v2 win rate among them for each checkpoint. For human evaluation, we recruited more than 50 human raters333The raters were paid contractors. They received their standard contracted wage, which is above the living wage in their country of employment. to rate the generated images of three pairs of network structures trained with PPO: (i) (gray-box) DiffCon vs. DiffCon-Naive, (i) (white-box) DiffCon-J vs. LoRA, (i) (white-box) DiffCon-S vs. LoRA. We randomly generate 50 prompts from the HPS-v2 prompt set. For each network structure we generate four images using different seeds for each prompt. The raters are tasked to evaluate a uniformly sampled pair of network structures out of the three, and are instructed to pick their favorite set based on the best-of-four to do evaluation. DiffCon parameterization. We instantiate the side network sθ=(zθ1,hθ2)s_θ=(z_ _1,h_ _2) for DiffCon as a lightweight UNet operating on the latent space (4+1 channels at 64×6464× 64 for SD v1.4). It consists of standard UNet blocks (DownBlock2D, CrossAttnDownBlock2D for the encoder, UpBlock2D, CrossAttnUpBlock2D for the decode, and CrossAttnMidBlock2D in the middle) similar as the ones used in diffusers von Platen et al. [2022]. See Table 2 for the details of the blocks. We found that using two DownBlock2D for the encoder and two UpBlock2D for the decoder works best, and the results presented in the main paper are under this structure. See Table 3 for the details of our side network structure under this choice. In our ablation study, we also tried other structures, see Appendix D for more details. Table 2: Block definitions used in sθs_θ. Common components: ResnetBlock2D: Standard ResNet block with GroupNorm, SiLU, and two Convolutional layers with kernel size 3×33× 3 and stride 11. Transformer2DModel: spatial transformer with cross-attention. Downsample2D: a convolutional layer with kernel size 3×33× 3 and stride 22. Upsample2D: 2×2× Nearest Neighbor Resize + 3×33× 3 conv, stride 11. Block Name Structure & Sequence CrossAttnDownBlock2D ResnetBlock2D Transformer2DModel Downsample2D (optional) DownBlock2D ResnetBlock2D Downsample2D (optional) CrossAttnMidBlock2D ResnetBlock2D Transformer2DModel ResnetBlock2D CrossAttnUpBlock2D ResnetBlock2D Transformer2DModel Upsample2D (optional) UpBlock2D ResnetBlock2D Upsample2D (optional) Table 3: Side network structure used in the main paper. Configuration: layers_per_block=1, block_out_channels=(128, 256), sample_size=64×6464× 64, down_block_types=(DownBlock2D, DownBlock2D), up_block_types=(UpBlock2D, UpBlock2D). Stage Block Sequence Input Channels Output Channels Resolution Details Input conv Conv2D 4 128 64×64→64×6464× 64→ 64× 64 ker 3×33× 3, stride 11. Down block 0 DownBlock2D 128 128 64×64→32×3264× 64→ 32× 32 Downsample: True. Down block 1 DownBlock2D 128 256 32×32→32×3232× 32→ 32× 32 Downsample: False. Mid block CrossAttnMidBlock2D 256 256 32×32→32×3232× 32→ 32× 32 - Up block 0 UpBlock2D 512 (256+256) 256 32×32→64×6432× 32→ 64× 64 Skip from Down block 1. Upsample: True. Up block 1 UpBlock2D 384 (256+128) 128 64×64→64×6464× 64→ 64× 64 Skip from Down block 0. Upsample: False. Output conv Conv2D 128 4 64×64→64×6464× 64→ 64× 64 ker 3×33× 3, stride 11. sθs_θ (DiffCon) vs. s¯θ s_θ (DiffCon-Naive). Both DiffCon and DiffCon-Naive use the same lightweight UNet-style side network backbone, but they differ in (i) what is fed into the side network and (i) how its output is used. In DiffCon, the side network is structured as sθ=(zθ1,hθ2)s_θ=(z_ _1,h_ _2) and is evaluated on the pretrained reverse mean μ0(xt,c,t) _0(x_t,c,t) (c.f. (27), (28)); one output head produces a scalar gate zθ1z_ _1 (implemented as an extra output channel followed by a linear head), while the remaining channels form hθ2h_ _2. In contrast, DiffCon-Naive uses an almost identical backbone but removes this gating head (one fewer output channel) and does not split outputs; it directly predicts an additive residual that is added to ϵ0 _0. We provide a detailed comparison in Table 4. DiffCon: sθs_θ DiffCon-Naive: s¯θ s_θ Backbone lightweight UNet same backbone (by design) Side-net input (μ0(xt,c,t),c,t)( _0(x_t,c,t),c,t) (xt,c,t)(x_t,c,t) Output parameterization (zθ1,hθ2)(z_ _1,h_ _2) single tensor (no split) Output channels (last conv) 55 (extra channel for zθ1z_ _1) 44 (one fewer channel) How output is used controller-form correction (27), (28) direct residual: ϵθ=ϵ0+λmodels¯θ _θ= _0+ _model s_θ Intuition gated, structured control term ungated additive correction Table 4: Side network variants in gray-box finetuning. Both methods share the same side-net backbone, but DiffCon uses a structured output (z,f)(z,f) applied in the controller-form correction, while DiffCon-Naive directly adds an unsplit residual. Appendix D Additional Experiment Results Example generations. Table 5, 6, 7 shows some generation examples for each method under the SFT, RWL, and PPO settings. In addition, Figure 4, 5, 6, 7 show some generated images with varying guidance strength λmodel _ model for both DiffCon and DiffCon-J methods under the SFT and RWL settings, demonstrating the effectiveness of guidance at inference time. Table 5: Qualitative comparison of different methods using SFT training. We compare Pretrained sd v1.4, DiffCon, DiffCon-Naive, LoRA, DiffCon-J, and DiffCon-S. Prompt Pretrained DiffCon DiffCon-Naive LoRA DiffCon-J DiffCon-S ”An intricate, biomechanical blue planet by Bruce Pennington” ”A bluejay is eating spaghetti” ”A crazy looking fish swimming in Alcatraz in the style of artgerm” ”An oil painting by Frank Frazetta depicting a humanoid spawn of dragons with scales, big eyes, and a menacing appearance” ”An oil painting by Rembrandt depicting a muscular cat wielding a weapon with dramatic clouds in the background” Table 6: Qualitative comparison of different methods using RWL training. We compare Pretrained sd v1.4, DiffCon, DiffCon-Naive, LoRA, DiffCon-J, and DiffCon-S. Prompt Pretrained DiffCon DiffCon-Naive LoRA DiffCon-J DiffCon-S ”A Victorian woman quietly sings by a lake at night surrounded by fireflies, moon, and stars, painted by Vincent van Gogh and Jacques-Louis David” ”Tetsuo and Kaneda engage in a race through Neo Tokyo in a pencil drawing featuring a scribbled style” ”A black cat wearing a suit and smoking a cigar” ”A dog sits in a white car with the door open” ”A painting of a cityscape with a red cloud shining light over a sea and bridge, by Greg Rutkowski and Thomas Kinkade” Table 7: Qualitative comparison of different methods using PPO training. We compare Pretrained sd v1.4, DiffCon, DiffCon-Naive, LoRA, DiffCon-J, and DiffCon-S. Prompt Pretrained DiffCon DiffCon-Naive LoRA DiffCon-J DiffCon-S ”A bike parked in front of a doorway” ”A cliff covered in blue blood” ”A knight in black armour holds a detailed silver sword in front of a castle” ”A lizard wearing sunglasses” ”A creepy eldritch monster in a Swedish forest photographed from a low angle with detailed, realistic features and soft colors, inspired by Lovecraftian horror and created by artist Simon Stålenhag” (a) 1/λmodel=0.11/ _ model=0.1 (b) 1/λmodel=0.31/ _ model=0.3 (c) 1/λmodel=0.51/ _ model=0.5 (d) 1/λmodel=1.01/ _ model=1.0 Figure 4: Generations with different guidance strengths λmodel _ model using SFT with DiffCon. Prompt: “A black cat wearing a suit and smoking a cigar”. (a) 1/λmodel=0.11/ _ model=0.1 (b) 1/λmodel=0.31/ _ model=0.3 (c) 1/λmodel=0.51/ _ model=0.5 (d) 1/λmodel=1.01/ _ model=1.0 Figure 5: Generations with different guidance strengths λmodel _ model using SFT with DiffCon-J. Prompt: “A bluejay is eating spaghetti”. (a) 1/λmodel=0.11/ _ model=0.1 (b) 1/λmodel=0.31/ _ model=0.3 (c) 1/λmodel=0.51/ _ model=0.5 (d) 1/λmodel=1.01/ _ model=1.0 Figure 6: Generations with different guidance strengths λmodel _ model using RWL with DiffCon. Prompt: “A crazy looking fish swimming in Alcatraz in the style of artgerm”. (a) 1/λmodel=0.031/ _ model=0.03 (b) 1/λmodel=0.081/ _ model=0.08 (c) 1/λmodel=0.1251/ _ model=0.125 (d) 1/λmodel=1.01/ _ model=1.0 Figure 7: Generations with different guidance strengths λmodel _ model using RWL with DiffCon-J. Prompt: “A cute humanoid cat soldier wears a yellow raincoat, carries a rifle, and ventures through a dense forest, looking back over their shoulder”. Full automatic evaluation metrics. In Table 8, 9, 10, we report the HPS-v2 win rate vs. the pretrained model together with CLIP, CLIP-Aesthetics, and PickScore on the HPS-v2 test prompt set, for each algorithm at its reported checkpoint. We can see that HPS-v2 is improved by our methods, and the metrics are not compromised during our finetuning. Method # Params ↓ Gray-box HPS-v2 win rate vs. pretrained model ↑ Clip ↑ Clip-Aesthetics ↑ PickScore ↑ Pretrained - - 0.500 0.2750 5.5163 0.2103 DiffCon (ours) 1.2×1.2× 10^7 ✓ 0.6667 ± 0.0028 0.2744 ± 0.0004 5.4754 ± 0.0325 0.2105 ± 0.0001 DiffCon-Naive 1.2×1071.2× 10^7 ✓ 0.5655 ± 0.0165 0.2740 ± 0.0002 5.5153 ± 0.0048 0.2105 ± 0.0001 LoRA 1.7×1071.7× 10^7 0.5766 ± 0.0137 0.2760 ± 0.0002 5.5144 ± 0.0014 0.2107 ± 0.0000 DiffCon-J (ours) 1.6×1.6× 10^7 0.6964 ± 0.0293 0.2747 ± 0.0002 5.4531 ± 0.0013 0.2103 ± 0.0001 DiffCon-S (ours) 1.6×1.6× 10^7 0.6964 ± 0.0018 0.2747 ± 0.0010 5.4678 ± 0.0245 0.2104 ± 0.0003 Table 8: SFT evaluation results on HPS-v2 test prompt set at step 1000. Method # Params ↓ Gray-box HPS-v2 win rate vs. pretrained model ↑ Clip ↑ Clip-Aesthetics ↑ PickScore ↑ Pretrained - - 0.500 0.2750 5.5163 0.2103 DiffCon (ours) 1.2×1.2× 10^7 ✓ 0.6815 ± 0.0111 0.2732 ± 0.0009 5.4448 ± 0.0265 0.2105 ± 0.0001 DiffCon-Naive 1.2×1071.2× 10^7 ✓ 0.5060 ± 0.0190 0.2736 ± 0.0008 5.4977 ± 0.0118 0.2101 ± 0.0000 LoRA 1.7×1071.7× 10^7 0.6109 ± 0.0059 0.2696 ± 0.0007 5.4950 ± 0.0069 0.2096 ± 0.0003 DiffCon-J (ours) 1.6×1.6× 10^7 0.6555 ± 0.0210 0.2712 ± 0.0007 5.5037 ± 0.0294 0.2100 ± 0.0004 DiffCon-S (ours) 1.6×1.6× 10^7 0.7091 ± 0.0155 0.2710 ± 0.0010 5.4714 ± 0.0311 0.2102 ± 0.0004 Table 9: RWL evaluation results on HPS-v2 test prompt set at step 2000. Method # Params ↓ Gray-box HPS-v2 win rate vs. pretrained model ↑ Clip ↑ Clip-Aesthetics ↑ PickScore ↑ Pretrained - - 0.500 0.2750 5.5163 0.2103 DiffCon (ours) 1.2×1.2× 10^7 ✓ 0.6957 ± 0.0152 0.2711 ± 0.0004 5.3370 ± 0.0219 0.2094 ± 0.0001 DiffCon-Naive 1.2×1071.2× 10^7 ✓ 0.5201 ± 0.0203 0.2748 ± 0.0006 5.5186 ± 0.0079 0.2106 ± 0.0000 LoRA 1.7×1071.7× 10^7 0.9048 ± 0.0084 0.2642 ± 0.0026 5.5631 ± 0.0642 0.2133 ± 0.0002 DiffCon-J (ours) 1.6×1.6× 10^7 0.9353 ± 0.0079 0.2660 ± 0.0004 5.6292 ± 0.0657 0.2142 ± 0.0007 DiffCon-S (ours) 1.6×1.6× 10^7 0.9315 ± 0.0164 0.2635 ± 0.0028 5.4989 ± 0.0846 0.2125 ± 0.0012 Table 10: PPO evaluation results on HPS-v2 test prompt set at step 2400. Our RWL v.s. DPOK [Fan et al., 2023, Algorithm 2]. We compare the heuristic reward weight loss in DPOK [Fan et al., 2023, Algorithm 2] and our proposed polynomial reward weighting (21) using LoRA in Figure 8. For DPOK, we set the learning rate to be 3e-6, the batch size to be 128 (twice of ours) and sample images (from the pretrained model for the teacher sampling setting, and from the current diffusion policy for the online sampling setting) every iteration (the sample frequency is also twice of ours). The other configurations for DPOK are the same as our RWL for LoRA. We can see that our proposed RWL achieves better HPS-v2 win rate than DPOK with 1/4 of the sample complexity. (a) Teacher sampling setting (b) Online sampling setting Figure 8: HPS-v2 win rate comparison between our polynomial reward weighting (LoRA, ours) and DPOK-RWL (LoRA) during RWL finetuning. (a) uses teacher sampling where the images for training at each iteration are generated by the pretrained model, and (b) uses online sampling where the images for training are generated by the current diffusion policy. Ablation studies on learning rates for RL algorithms. We found LoRA performs better with higher learning rate (1e-4) for RWL and PPO, while DiffCon performs better with a lower learning rate (1e-5), as measured by HPS win rate (Figure 9). (a) RWL+DiffCon (b) RWL+LoRA (c) PPO+DiffCon (d) PPO+LoRA Figure 9: Learning rate ablation for RWL and PPO on HPS-v2 win rate (comparing 1e-4 vs. 1e-5). Ablation studies on τKL _ KL for PPO. We provide the effectiveness of different τKL _ KL for PPO in Figure 10 to justify our choice of τKL= _ KL=1e-4 in the main paper, which leads to the best performance. Note that when the KL regularization is too small (or is 0), the reward won’t go up. (a) DiffCon (b) LoRA (c) DiffCon-J Figure 10: Ablation studies on τKL _ KL for PPO on HPS-v2 win rate (five learning curves correspond to τKL∈ _ KL∈\0, 1e-5, 1e-4, 5e-4, 1e-3\). τKL= _ KL=1e-4 leads to the best performance. Ablation studies on τRWL _ RWL and weighting functions for RWL. In the main paper, we report RWL results using the polynomial reward weighting function wαw_α (Eq. (21)). We first ablate τRWL _ RWL under this polynomial weighting: smaller τRWL _ RWL generally leads to better performance (Figure 11), but too small τRWL _ RWL (smaller than 0.0005) will cause instability in training (NaN in reward weight computation). Thus we choose τRWL=0.0005 _ RWL=0.0005 in the main paper. (a) DiffCon (b) LoRA Figure 11: Ablation studies on τRWL _ RWL for RWL (three HPS-v2 win rate curves in each subfigure correspond to τRWL∈0.0005,0.0025,0.005 _ RWL∈\0.0005,0.0025,0.005\). τRWL=0.0005 _ RWL=0.0005 leads to the best performance. We additionally report RWL results under the exponential reward weighting function wKLw_KL (Eq. (20)) with τRWL=0.0015 _ RWL=0.0015 in Figure 12. Consistent with our main results, our methods achieve higher HPS win rate than the baselines, and DiffCon (with fewer trainable parameters) outperforms the strong white-box LoRA baseline. Figure 12: RWL under the exponential reward weighting function wKLw_KL (Eq. (20)) with τRWL=0.0015 _ RWL=0.0015 on HPS-v2 win rate. Ablation studies on side network structures. As mentioned in Appendix C, we found that using two DownBlock2D for the encoder and two UpBlock2D for the decoder (i.e., down_block_types=(DownBlock2D, DownBlock2D), up_block_types=(UpBlock2D, UpBlock2D)) works best. We also tried different side network block combinations; see Table 11 and Figure 13 for more details. Table 11: Side network structure variants used in the ablation study. Label # trainable params down_block_types up_block_types block_out_channels 1Res 11,702,717 DownBlock2D UpBlock2D 256 1Res+1Attn 13,391,549 CrossAttnDownBlock2D, DownBlock2D UpBlock2D, CrossAttnUpBlock2D (128, 256) 2Attn 18,519,485 CrossAttnDownBlock2D, CrossAttnDownBlock2D CrossAttnUpBlock2D, CrossAttnUpBlock2D (128, 256) 2Res 11,810,621 DownBlock2D, DownBlock2D UpBlock2D, UpBlock2D (128, 256) (a) SFT (b) RWL (c) PPO Figure 13: Ablation studies on DiffCon block structures evaluated by the HPS-v2 win rate curve. Each plot contains four curves corresponding to 1Res, 1Res+1Attn, 2Attn, and 2Res (main-paper setting). Ablation studies on the input of the side network. Under our parameterization (27), we input μ0(xt,c,t) _0(x_t,c,t) instead of xtx_t to the side network. We found this leads to better performance for RWL, see Figure 14. Figure 14: Ablation studies on the input of DiffCon for RWL, evaluated by the HPS-v2 win rate curve. The plot contains two curves corresponding to input_mean=True or False, i.e., using the mean μ0(xt,c,t) _0(x_t,c,t) or xtx_t as the input to the side network sθs_θ.