Variational Optimal Control for Diffusion Guidance

Ruoyu Wang; Xinyi Ye; Yunshuo Shi; Liangyu Yuan; Beier Zhu; Chi Zhang

cs.LGOct 27, 2026

Variational Optimal Control for Diffusion Guidance

Authors: Ruoyu Wang, Xinyi Ye, Yunshuo Shi, Liangyu Yuan, Beier Zhu, Chi Zhang

Organizations: AGI Lab, Westlake University · Shandong University · Shanghai Jiao Tong University · University of Science and Technology of China

Abstract

We present a variational control framework that connects sampling-time guidance and reward-based model adaptation through a shared Jacobian pullback. Starting from terminal reward maximization with a path-space KL penalty, the Gibbs variational principle and Doob transform yield a correction to the generative dynamics. A local denoising approximation exposes the pullback that transports output preferences to the controlled variable. This construction gives Variational Optimal Guidance (VOG) in state space and extends to Distilled Variational Optimal Control (DVOC) by amortizing locally improved endpoints into model parameters. The resulting controls share a local variational rule: VOG acts through the denoising Jacobian, while DVOC realizes the parameter pullback by endpoint distillation. Under a linear MMSE model, VOG's exact correction lies in the data subspace. On ImageNet, VOG reduces DiT-XL/2 FID from 2.27 to 2.04 at 250 sampling steps and from 7.58 to 4.44 at 25 steps. After 20 updates on Z-Image-Turbo, DVOC improves HPSv3++ and four held-out evaluator scores over the untrained model, with the highest corresponding means among the trained methods compared. These results connect inference-time control and reward adaptation through a common geometric principle.

Figures & tables

Figure 1: Variational control for sampling-time guidance and model adaptation. The gray curve illustrates generation from noise to the sampled endpoint. VOG (blue) evaluates the objective at the denoised prediction $\hat{x}_{0}$ and maps its gradient through $\mathbf{J}_{t}$ to correct the current sampling state. DVOC (pink) caches the trajectory, refines the actual endpoint $y=x_{0}$ within a local neighborhood, and distills the fixed target $y^{+}$ into LoRA parameters. The three ringed states supply the selected velocity evaluations ( $K=3$ ) used for endpoint distillation. Target refinement uses two projected reward-ascent steps ( $J=2$ ).

Figure 2: Toy class-conditional generation on fractal-like 2D branching-tree distributions [ 18 ] under two increasingly challenging settings: a two-class benchmark (CLS $=2$ ) and a three-class variant with shared structure (CLS $=3$ ), which induces stronger class overlap. We quantify generation difficulty using $\mathrm{tr}(\mathrm{Cov})$ of the class-conditional ground-truth mixture, where larger values indicate greater difficulty. (a) Ground truth samples from the target class distribution. (b) No guidance: conditional sampling with a small diffusion model produces outliers and off-class leakage. (c) CFG ( $\omega=4$ ) suppresses outliers but shows mode collapse, reducing branch coverage; the collapse is more severe as difficulty increases from $\mathrm{tr}(\mathrm{Cov})=0.228$ (CLS $=2$ ) to $\mathrm{tr}(\mathrm{Cov})=0.403$ (CLS $=3$ ). (d) VOG (ours) preserves diversity and branch coverage while achieving strong conditional alignment. Detailed quantitative results are provided in Table S11 .

	EDM2-S (512 $\times$ 512)			DiT-XL/2 (256 $\times$ 256)			EDM2-L (64 $\times$ 64)			DiT-XL/2 (512 $\times$ 512)
Guidance	FID	Precision	Recall	FID	Precision	Recall	FID	Precision	Recall	FID	Precision	Recall
CFG [ 17 ]	2.29	0.751	0.582	2.27	0.829	0.584	1.99	0.710	0.601	3.04	0.838	0.541
+ VOG (Ours)	2.03	0.755	0.591	2.04	0.756	0.607	1.42	0.714	0.597	2.89	0.745	0.560
Interval [ 21 ]	1.87	0.760	0.598	1.97	0.801	0.632	1.64	0.730	0.597	2.41	0.792	0.611
+ VOG (Ours)	1.67	0.742	0.614	1.94	0.765	0.611	1.52	0.720	0.593	2.39	0.760	0.567
Cosine [ 12 ]	2.15	0.770	0.619	2.30	0.861	0.520	1.67	0.726	0.587	3.48	0.713	0.586

Table 1: VOG as a plug-in correction on top of CFG-style guidance baselines, on class-conditional ImageNet. Shaded rows apply VOG to the baseline immediately above them. We report FID (

\downarrow

), Precision (

\uparrow

), and Recall (

\uparrow

); within each column the best result is bolded and the second best underlined . A dash indicates a configuration the baseline does not support.

Z-Image-Turbo, 1024 2 , nine-step Euler sampler
Method	In-Domain				Out-of-Domain
	HPSv3++		CLIP		Q-Align Q.		Q-Align A.		AES		ImageReward		DeQA
	Score $\uparrow$	$\Delta$	Score $\uparrow$	$\Delta$	Score $\uparrow$	$\Delta$	Score $\uparrow$	$\Delta$	Score $\uparrow$	$\Delta$	Score $\uparrow$	$\Delta$	Score $\uparrow$	$\Delta$
Z-Image-Turbo (untrained)	6.72	—	0.2757	—	4.825	—	3.533	—	5.412	—	0.959	—	4.439	—
ReFL [ 42 ]	12.22	$+5.50$	0.2744	$-0.0013$	4.844	$+0.020$	4.223	$+0.691$	5.955	$+0.543$	1.159	$+0.200$	4.133	$-0.306$
DRTune [ 40 ]	12.70	$+5.97$	0.2688	$-0.0070$	4.932	$+0.107$	4.489	$+0.956$	6.095	$+0.683$	1.225	$+0.266$	4.262	$-0.177$

Table 2: Reward alignment after 20 optimization updates on Z-Image-Turbo. The untrained base and four trained methods use the same 1,000 prompts and initial noise. HPSv3++/CLIP are in-domain; the others are OOD.

\Delta

: score minus untrained base. Higher is better. Best/second-best: bold / underlined . Blue: DVOC.

Figure 3: Evaluation and training dynamics. (a,b) Online EMA evaluation every ten updates, using HPSv3++ condition 0.0 and trainer-specific noise sampling. (c) Training reward,

0.03846\,R_{\mathrm{HPSv3++}}+R_{\mathrm{CLIP}}

: ReFL/DRTune score late-state proxies; DVOC/OPSD score sampled endpoints. (d) DVOC training reward for

K=1,2,3

. Training panels show raw means (faint) and trailing five-update averages; evaluation is unsmoothed. The vertical line marks update 20. DVOC uses

K=3

in (a–c); (d) is the ablation.

	In-domain		OOD
Variant	HPSv3++	CLIP	Q-Quality	Q-Aesth.	AES	ImageReward	DeQA
Endpoint, $J=2$	13.727	0.2760	4.947	4.640	6.281	1.2678	4.252
Endpoint, $J=1$	13.653	0.2781	4.927	4.629	6.211	1.2673	4.131
Proxy, $J=2$	13.705	0.2763	4.935	4.618	6.277	1.2684	4.242

Table 3: Target construction at fixed

K=3

. Three paired runs use 20 updates and 1,000 matched test images. All scores are higher-is-better. Bold marks the highest mean within this ablation.

DVOC variant	In-Domain		Out-of-Domain
DVOC variant	HPSv3++ $\uparrow$	CLIP $\uparrow$	Q-Align Q. $\uparrow$	Q-Align A. $\uparrow$	AES $\uparrow$	ImageReward $\uparrow$	DeQA $\uparrow$
$K=1$	13.36	0.2729	4.948	4.628	6.272	1.224	4.249
$K=2$	13.69	0.2751	4.947	4.610	6.274	1.257	4.263
$K=3$ (default)	13.79	0.2765	4.943	4.618	6.288	1.279	4.237

Table 4: Active-step ablation. Only

K

changes;

J=2

and the endpoint radius are fixed. Means use the same 1,000 paired prompts at update 20. Bold/underlining mark best/second-best within this table.

Appendix figures & tables19 assets

Supplementary material from the paper’s appendix.

Appendix

Symbol	Meaning
$x_{t}$ , $\hat{x}_{0}$	noisy state at time $t$ ; denoised prediction of the clean sample
$\epsilon_{c}$ , $\epsilon_{u}$ , $\epsilon_{\mathrm{vog}}$	conditional and unconditional predictions; VOG -corrected prediction
$\bar{\alpha}_{t}$ , $\sqrt{1-\bar{\alpha}_{t}}$	cumulative signal power and VP noise standard deviation
$c$ , $f(x_{0},c)$	conditioning signal; terminal reward on the clean sample
$\omega$ , $w=\omega-1$	CFG guidance scale; its conditional-centered form, the control strength
$\lambda$	reward temperature in the VOC objective

Appendix

Table S1: Notation used in the main text.

Controller	FID ( $\downarrow$ )	Relative runtime
Pullback through $\epsilon_{\mathrm{cfg}}$	5.48	–
Pullback through $\epsilon_{c}$ (VOG)	2.04	$1.55\times$

Appendix

Table S2: Differentiation-branch ablation on DiT-XL/2 (

256\times 256

), with VJP stride

K_{\mathrm{VOG}}=4

. Runtime is relative to CFG.

Resolution	Method	$\omega$	$K$	$\beta$	$s$	\Block Guidance
interval	$\text{EMA}_{m}$	$\text{EMA}_{g}$
\Block [fill=white]10-1 $64\times 64$	EDM2-L
	Classifier-free guidance	1.15	-	-	-	Full	0.040	0.040
	+ VOG (Ours)	1.05	2	0.15	-	Full	0.040	0.040
	Guidance interval	1.15	-	-	-	(0.3,1.5]	0.040	0.040
	+ VOG (Ours)	1.10	2	0.15	-	(0.3,1.5]	0.040	0.040

Appendix

Table S3 : Hyperparameters and settings for ImageNet-64 and ImageNet-512. The experiments are conducted using EDM2-L and EDM2-S architectures. The column

\omega

denotes the guidance scale. The column Guidance interval specifies the range of noise levels

\sigma

where classifier-free guidance is applied. For our proposed method (+Ours),

k

represents the frequency of regularization steps (computing gradients every

k

steps), and

\beta

indicates the momentum coefficient. The parameter

s

refers to the exponent used in Guidance cosine. The columns

\text{EMA}_{m}

and

\text{EMA}_{g}

indicate the EMA decay lengths for the main generative model and the guiding model, respectively.

Resolution	Method	$\omega$	$K$	$\beta$	$s$	\Block Guidance
interval
\Block [fill=white]6-1 $256\times 256$	Classifier-free guidance	1.50	-	-	-	Full
	+ VOG (Ours)	1.50	4	0.15	-	Full
	Guidance interval	1.50	-	-	-	(0.34, 1.02]
	+ VOG (Ours)	1.90	2	0.15	-	(0.34, 1.02]
	Guidance cosine	1.50	-	-	0.3	Full

Appendix

Table S4 : Hyperparameters and settings for ImageNet-256 and ImageNet-512. The experiments are conducted using the DiT-XL/2 architecture. The column

\omega

denotes the guidance scale. The column Guidance interval specifies the range of noise levels

\sigma

where classifier-free guidance is applied. For our proposed method (+Ours),

k

represents the frequency of regularization steps (computing gradients every

k

steps), and

\beta

indicates the momentum coefficient. The parameter

s

refers to the exponent used in Guidance cosine.

Dataset	$\bm{K}$	$\bm{K_{\text{ent}}}$	$\bm{K_{\text{hhi}}}$	Top 5% Mass	$\bm{\text{tr}(\text{Cov})}$	MST
CLS-2 (Class A)	1016	249.81	90.29	0.5429	0.2275	15.35
CLS-3 (Class A)	1036	214.05	102.91	0.6361	0.4029	18.57

Appendix

Table S5 : Quantitative difficulty metrics for the 2D synthetic datasets (Class A).

Method	Backward path	Reward target or loss	Weight decay
ReFL	One selected late-state evaluation	Negative standardized reward; loss scale $10^{-3}$ ; no hinge	$10^{-2}$
DRTune	Three connected evaluations in a truncated rollout	Same reward loss and scale as ReFL; early-stop ratio 0.4	$10^{-2}$
DiffusionOPSD	One queried state	Positive and negative PGD targets; old-policy mixing; $J=2$ , $\rho=0.1$ , $\eta=1$	$10^{-4}$
DVOC ( $K=3$ )	Three cached states; detached state dependencies	Actual-endpoint PGD target; $J=2$ , $\rho=0.1$ , $\eta=1$ ; normalized target MSE	$10^{-4}$

Appendix

Table S6: Method-specific training settings. A differentiated evaluation carries a parameter gradient.

Method	GPUs	Seconds / update	Minutes / update	GPU-min / update
ReFL	8	$272.8\pm 1.0$	4.55	36.4
DRTune	8	$336.8\pm 1.3$	5.61	44.9
DiffusionOPSD	8	$397.6\pm 0.8$	6.63	53.0
DVOC ( $K=1$ ) †	8	$325.8\pm 0.6$	5.43	43.4
DVOC ( $K=2$ ) †	8	$391.2\pm 0.4$	6.52	52.2
DVOC ( $K=3$ )	8	$458.0\pm 0.5$	7.63	61.1

Appendix

Table S7: Training time per optimizer update on one eight-GPU node. Means and standard deviations use 18 intervals from updates 2–20, excluding update 11 because it follows evaluation and checkpointing. Each update uses 576 training trajectories. Standard deviations describe timing variation within a single run.

\dagger

denotes estimates from rollout-completion intervals; other rows use optimizer-completion intervals.

Metric	Mean difference	95% interval
HPSv3++	+7.06357	[+6.90034, +7.22679]
CLIP	+0.00077	[-0.00142, +0.00297]
Q-Align Q.	+0.11801	[+0.10556, +0.13046]
Q-Align A.	+1.08532	[+1.05159, +1.11906]
AES	+0.87600	[+0.84961, +0.90238]
ImageReward	+0.31966	[+0.28264, +0.35668]

Appendix

Table S8 : Paired differences: default DVOC (

K=3

) minus untrained Z-Image-Turbo. Normal-approximation 95% intervals use 1,000 paired prompts and are reported per metric.

Figure S1 : Qualitative comparison of reward-based adaptation. Columns show Z-Image-Turbo, ReFL, DRTune, DiffusionOPSD, and DVOC. Rows show (a) a museum interior, (b) a linen jacket, (c) a bicycle, and (d) a reading room. The examples illustrate differences in material texture, lighting, and local detail across methods.

Model	Res/BS	CFG/ VOG Runtime (sec)
EDM2-S	64/8	3.66 / 5.33 (1.46x)
DiT-XL/2	256/8	10.97 / 17.04 (1.55x)
EDM2-S	512/8	4.66 / 9.22 (1.98x)
DiT-XL/2	512/8	39.74 / 65.03 (1.64x)
Model	Res/BS	CFG/ VOG Memory (GB)
EDM2-S	64/1	1.18 / 1.79 (1.53x)

Appendix

Table S9 : Runtime and peak memory of VOG relative to CFG, reported as CFG / VOG with the ratio in parentheses. Res/BS denotes resolution and batch size; all measurements are taken on a single NVIDIA RTX 4090.

Model	Method	FID	Recall	Precision	$\omega$	$\beta$	$K$
\Block [fill=white]2-1DiT-XL/2	CFG	7.58	0.429	0.784	1.50	-	-
	VOG	4.44	0.545	0.699	1.50	0.15	2
\Block [fill=white]2-1SiT-XL/2	CFG	6.20	0.555	0.685	1.50	-	-
	VOG	3.89	0.516	0.735	1.10	0.15	2

Appendix

Table S10 : Quantitative comparison of 25-step sampling on DiT-XL/2 and SiT-XL/2 (

256\times 256

, ImageNet). We report FID (

\downarrow

), Recall (

\uparrow

), and Precision (

\uparrow

), together with the guidance scale

\omega

, the momentum coefficient

\beta

, and the VJP recomputation interval

K

	Two-class (CLS=2)				Three-class (CLS=3)
Method	Prec.	Rec.	F1	Wass.	Prec.	Rec.	F1	Wass.
Ground Truth	0.9916	1.0000	0.9958	0.0021	0.9862	1.0000	0.9931	0.0026
No Guidance	0.8948	0.9989	0.9440	0.0031	0.9194	0.9970	0.9567	0.0035
Standard CFG	0.9110	0.8755	0.8929	0.0064	0.9392	0.7785	0.8514	0.0071
VOG (ours)	0.9230	0.9568	0.9396	0.0043	0.9536	0.9328	0.9431	0.0042

Appendix

Table S11 : Quantitative results on 2D branching-tree generation. We report Precision, Recall, F1-score (higher is better), and Wasserstein distance (lower is better) for both Two-class and Three-class settings. Best results (excluding Ground Truth) are highlighted in bold .

Panel of Figure 20 — Figure S2 : Ablation studies on guidance-scale sensitivity (left) and hyperparameter robustness (right).

Figure S3 : Quantitative comparison on EDM2-S (

512\times 512

) across varying guidance scales. We report FID (

\downarrow

), Precision (

\uparrow

), and Recall (

\uparrow

) as a function of the guidance weight

\omega_{\text{CFG}}

. VOG (red) demonstrates superior robustness at high scales compared to standard CFG (blue), maintaining an FID of

\approx 2.4

\omega=2.0

while the baseline degrades beyond 4.0.

Figure S4 : Quantitative comparison on DiT-XL/2 (

256\times 256

) across varying guidance scales. While standard CFG suffers from sharp performance degradation at high guidance strengths (e.g., FID

>3.40

\omega=1.7

), VOG maintains structural fidelity (high Recall) and generation quality (FID

<2.60

), effectively preventing manifold drift.

EDM2-S
Benchmark	FID	Precision	Recall
+AutoGuidance	1.044	0.701	0.626
+ VOG (Ours)	1.031	0.696	0.620
EDM2-L
+Standard CFG	1.993	0.710	0.601
+ VOG (Ours)	1.418	0.714	0.597

Appendix

Table S12 : VOG composed with different guidance heuristics on ImageNet

64\times 64

. Each shaded row applies VOG on top of the heuristic immediately above it.

Figure S5 : Training reward histories. Separate panels distinguish sampled endpoints from late-state proxies. Thin lines show every raw update mean; strong lines show the trailing five-update mean. The vertical reference marks update 20. The HPSv3++ and CLIP panels show DVOC and DiffusionOPSD.

Figure S6 : Online evaluation histories. HPSv3++ and CLIP are in-domain evaluators; ImageReward is held out of training. Each marker is a checkpoint mean, with no smoothing. The vertical line marks update 20. The HPSv3++ condition and noise sampling differ from the paired evaluation in the main table.

Figure S7 : DVOC active-step ablation over optimization.

K=1,2

appear only in the ablation; the main method remains

K=3

. Training reward includes all raw points and a trailing five-update mean. Evaluation panels show only recorded checkpoint means. The vertical line marks the main paper’s update-20 checkpoint.

Diffusion Model Guidance

Diffusion model guidance steers denoising toward desired conditions, rewards, constraints, or away from unwanted concepts.

See all 364 papers

Diffusion Controller: Framework, Algorithms and Parameterization
Mar 7, 2026Tong Yang, Moonkyung Ryu, Chih-Wei Hsu +4Diffusion Model Adaptation Diffusion Model Fine-Tuning
Carnegie Mellon University; Intern at Google Research · Google · Yale University
Hierarchical Variational Policies for Reward-Guided Diffusion
May 20, 2026Kushagra Pandey, Farrin Marouf Sofian, Jan Niklas Groeneveld +2Few-Step Diffusion Sampling Diffusion Model Guidance
Department of Computer Science University of California Irvine
VGAS: Variance-Reduced Guidance and Adaptive Selection for Training-Free Reward Alignment in Discrete Diffusion
Aug 27, 2026Kwanyoung KimDiffusion Model Alignment Diffusion Model Guidance
Department of AI, GIST

Mar 7, 2026J/K move · Enter open · S save

Diffusion Controller: Framework, Algorithms and Parameterization

Tong Yang, Moonkyung Ryu, Chih-Wei Hsu, Guy Tennenholtz+3

Carnegie Mellon University; Intern at Google Research · Google · Yale University

Diffusion Model Adaptation Diffusion Model Fine-Tuning Diffusion Models Diffusion Model Alignment RL for Diffusion Models Diffusion-Based Denoising

Open paper AskPDF arXiv

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.

Abstract

Diffusion Model Guidance

Explore similar work

Diffusion Controller: Framework, Algorithms and Parameterization

Hierarchical Variational Policies for Reward-Guided Diffusion

VGAS: Variance-Reduced Guidance and Adaptive Selection for Training-Free Reward Alignment in Discrete Diffusion