Variational Optimal Control for Diffusion Guidance

• Users seek preferred image content and attributes beyond faithful distribution reproduction, but stronger guidance can reduce coverage and external objectives can favor directions departing from data structure.The paper asks how terminal objectives and intervention costs determine control across noisy states and model parameters.Source: §1 · pp. 1–2
• VOC balances terminal reward against path-space KL deviation from a pretrained reference, and its local denoised-reward approximation transports the output gradient through the denoising Jacobian to control the current state.Under a linear subspace model with an optimal denoiser, the exact correction lies in the data subspace; the practical sampler approximates this pullback for neural denoisers.Source: §1 · p. 2
• DVOC refines a sampled terminal latent with bounded-neighborhood reward gradients, then distills the improved target into model parameters through a parameter-Jacobian pullback.Selected denoising steps carry the parameter gradient, while the adapted model uses its ordinary sampler at inference.Source: §1 · p. 2
• VOG improves FID in all twelve ImageNet model–guidance pairings, reducing 25-step DiT-XL/2 FID from 7.58 to 4.44; after 20 updates, DVOC improves HPSv3++ and four held-out scores over pretrained Z-Image-Turbo and leads trained comparators.Source: §1 · p. 2

• Sampling guidance spans classifier gradients [9], CFG [16], guidance adjustments [5] [18] [21] [29] [30] [41], manifold-aware methods [13] [14], and clean-prediction objective propagation [1] [4] [45]; this work derives state corrections from terminal control through the denoising Jacobian.Source: §2 · p. 2
• Reward adaptation includes policy-gradient methods [2] [11], preference learning [36], reward backpropagation [6] [27] [40] [42], and distillation [37] [47]; this work formulates adaptation as amortized control through a shared local variational principle and Jacobian pullback.Source: §2 · p. 2
• Figure 1 contrasts VOG’s objective evaluation at a denoised prediction with DVOC’s cached trajectory, locally refined endpoint, and distillation into LoRA parameters.The depicted DVOC example uses three selected velocity evaluations and two projected reward-ascent steps.Source: Figure 1 · p. 3, §2 · p. 3
• Variational control relates stochastic reward-versus-reference optimization to control regression [10] and Doob-transform conditioning or editing [24], while VOG and DVOC realize state and parameter corrections as trajectory guidance and generator amortization.Source: §2 · p. 3

3 Variational Optimal Control and Model Geometry

• VOG controls the current state through denoised-prediction sensitivity, whereas DVOC controls model parameters through endpoint sensitivity, expressing a shared local variational principle.The section introduces a low-dimensional linear-subspace model to characterize the denoising Jacobian.Source: §3 · p. 3

3.1 How CFG Shifts the Mean and Reduces Coverage

• CFG’s transported class-uncertainty contrast can shift the generated mean, while stronger guidance suppresses outliers but reduces branch coverage, especially as class overlap increases.The analysis assumes Bayes-optimal branch mixing and locally affine class denoisers with shared contraction.Source: §3.1 · pp. 4–5, Figure 2 · p. 4

3.2 Variational Optimal Control (VOC)

• VOG locally pulls a denoised reward gradient back through the denoising Jacobian to correct the conditional reference dynamics derived from terminal-reward control.The control objective penalizes deviation from a conditional reference process; its Doob transform shifts the drift while preserving diffusion covariance.Source: §3.2 · pp. 5–6

3.3 Geometric Consistency Analysis

• Under the linear latent-subspace assumption and optimal linear MMSE denoising, VOG’s exact correction lies in the data subspace and filters out reward-gradient components normal to it.The neural-denoiser implementation approximates this exact pullback.Source: §3 · p. 3, §3.3 · p. 6

3.4 Practical Implementation Strategies

• Practical VOG strategies cache pullbacks between recomputations, use a factorized VJP surrogate, and smooth the guidance signal with exponential momentum.The inference stride for cached corrections is distinct from DVOC’s active-step count K.Source: §3.4 · p. 6

3.5 From State Control to Model Adaptation

• DVOC refines sampled terminal latents within a local displacement budget, then distills those targets into model parameters through an endpoint parameter Jacobian.Averaging over sampled prompts and trajectories amortizes terminal controls into a shared model that uses its ordinary sampler at inference.Source: §3.5 · pp. 7–8

4.1 Evaluation Protocols

• The evaluation uses ImageNet at three resolutions and reports FID, Precision, and Recall; reward adaptation compares four trained methods with untrained Z-Image-Turbo.The adaptation evaluation uses the complete 1,000-prompt Pick-a-Pic split, matched prompts and initial noise, and one training seed.Source: §4.1 · pp. 8–9

4.2 Sampling-Time Guidance Results

• On the branching-tree benchmark, VOG attains Recall 0.957 and 0.933 across two- and three-class settings while achieving the highest Precision among guided samplers.CFG reduces Recall to 0.876 and 0.779 in those settings.Source: §4.2 · p. 9
• VOG improves FID in all twelve ImageNet model–guidance pairings, including DiT-XL/2 at 25 steps from 7.58 to 4.44.The reported 25- and 50-step results are 7.58 to 4.44 and 3.69 to 2.99, respectively.Source: §4.2 · p. 9

4.3 Reward Alignment and Generalization

• DVOC raises HPSv3++ from 6.72 to 13.79 and leads trained methods on four held-out evaluators, while DeQA falls below the base score.The held-out leads are Q-Align quality 4.943, Q-Align aesthetics 4.618, AES 6.288, and ImageReward 1.279; DeQA is 4.237 versus 4.439 for the base.Source: §4.3 · pp. 9–10
• From update 20 to 100, DVOC’s online CLIP rises from 0.2768 to 0.3111 and HPSv3++ from 12.71 to 15.59.These online evaluations use a different protocol from the paired main table.Source: §4.3 · p. 10

4.4 DVOC Ablations

• The preferred active-step budget depends on the evaluator: K = 3 leads on HPSv3++, CLIP, AES, and ImageReward, K = 1 on both Q-Align scores, and K = 2 on DeQA.At fixed K = 3, endpoint J = 2 improves Q-Align quality, AES, and DeQA over one ascent step, while CLIP decreases; the proxy endpoint target also improves three listed metrics over the proxy.Source: §4.4 · p. 11

4.5 Computational Cost

• Increasing K expands the parameter pullback while holding sampled trajectories and target construction fixed, with mean update times of 325.8, 391.2, and 458.0 seconds for K = 1, 2, and 3.K = 1 reduces update time by 28.9% relative to K = 3; at fixed K = 3, one versus two target steps take 414.1 versus 461.7 seconds per update.Source: §4.5 · p. 11

• VOG improves FID across evaluated ImageNet settings, while DVOC after 20 Z-Image-Turbo updates improves HPSv3++ and four held-out scores over the base and leads trained comparators on those metrics.Source: §5 · p. 12
• VOG controls individual trajectories, whereas DVOC amortizes endpoint improvements into the generator through a shared Jacobian pullback.Source: §5 · p. 12

A.1 Diffusion and Guidance

• The VP perturbation combines a clean sample with Gaussian noise, and the noise predictor defines the model score.Source: §A.1 · p. 12
• Classifier guidance adds a classifier log-probability gradient, while CFG extrapolates between conditional and unconditional predictions and may reduce coverage at large scale.The CFG guidance scale is ω.Source: §A.1 · p. 12

A.2 Summary of Notation

• The notation distinguishes CFG scale ω, VOG correction schedule ζ(t), momentum coefficient β, and VP noise rate βvp(t).Source: §A.2 · p. 12

B.1 Optimal path measure

• The tilted path measure uniquely maximizes the terminal-reward objective with value ℓt(x).The continuation is conditioned on the current state and prompt.Source: §B.1 · p. 12, B Derivation of Variational Optimal… · p. 12

B.2 Drift correction and control cost

• The optimal drift adds at∇ℓt to the reference drift while leaving the reference diffusion covariance unchanged.Source: §B.2 · p. 13

B.3 Local transition and Jacobian pullback

• Linearizing the desirability over a Gaussian reference transition shifts its mean by Σtgt without changing its covariance.This yields the local transition correction used in the derivation.Source: §B.3 · p. 13

C.1 Conditional clean-prediction pullback

• Evaluating reward on the conditional clean prediction maps its gradient through the denoising Jacobian to correct the sampling state.The CFG controller uses the conditional–unconditional score contrast as a proxy for the clean-space log-ratio gradient.Source: §B.3 · p. 14, §C.1 · p. 14, C The CFG-Specific VOG Controller · p. 14

C The CFG-Specific VOG Controller

• The CFG-specific objective uses a conditional log-ratio, with standard CFG written as ϵc + wgt.Source: C The CFG-Specific VOG Controller · p. 14

C.2 Factorized VJP

• The factorized VJP drops a residual that vanishes for a diagonal Jacobian or when the contrast is constant over each influence set.The practical controller combines a sensitivity factor with the guidance contrast element-wise.Source: §C.2 · p. 14

C.3 Choice of differentiation branch

• Differentiating the CFG prediction adds a term quadratic in guidance strength, whereas the conditional branch keeps its sensitivity map independent of w at a fixed state.The differentiation-branch ablation is on DiT-XL/2 at 256 × 256 with VJP stride 4.Source: §C.3 · p. 15, Table S2 · p. 15

D.1 Endpoint optimization

• DVOC refines the sampled endpoint with normalized projected reward ascent inside a radius proportional to the endpoint norm.The specified configuration uses two ascent steps and a detached final target.Source: §D.1 · p. 15

D.2 Cached endpoint connector

• The cached endpoint connector equals the endpoint in the forward pass but routes parameter gradients through selected velocity evaluations.The resulting fitting objective is used to optimize LoRA parameters.Source: §D.2 · p. 16

E.1 Class contrast and guided dynamics

• The VP probability-flow analysis derives CFG’s class-contrast effect under Bayes-optimal branch mixing and locally affine class denoisers.The ODE is integrated from noise time T to clean time 0.Source: §E.1 · p. 16

E.2 Transport of the guidance perturbation

• The first-order mean shift includes a path-dependent transport kernel, and its class-posterior gap also depends on guidance weight.Source: §E.2 · pp. 16–17
• When class means lie in S and the local filter preserves S, VOG’s normal response vanishes; learned-denoiser cross-subspace coupling can alter that response.The exact linear-MMSE correction is established to lie in the data subspace.Source: §E.2 · p. 17

F.1 General Configuration and Hyperparameters

• The experiments use 50k generated samples per FID configuration, tune guidance weight by architecture and task, and evaluate the DVOC comparison after 20 updates.Source: §F.1 · pp. 17–18

F.2 Structure of the 2D Synthetic Distributions

• CLS-3 is harder than the other branching-tree dataset, with greater spatial dispersion and more concentrated probability mass in its largest modes.Its covariance trace is approximately 0.40 versus 0.23, and its Top 5% Mass is approximately 63.6%.Source: §F.2 · p. 18

F.3 DVOC and Baseline Reproduction Details

• The DVOC comparison uses paired prompts and initial noise, while its trained-method rows are fixed update-20 checkpoints from longer runs with differing gradient paths, reward calls, and per-update costs.The reported mean update times over 90 intervals are 272.9 s for ReFL, 337.2 s for DRTune, 396.8 s for DiffusionOPSD, and 456.5 s for DVOC.Source: §F.3 · pp. 19–20

F.4 Endpoint-Target Ablations

• The paired 20-update runs compare endpoint targets with one or two ascent steps against a two-step late-state proxy using matched prompts, noise, active states, and reward evaluations.The proxy transfers its optimized displacement to the endpoint.Source: §F.4 · p. 21
• 0.0202, 0.0692, and 0.1203 are the Q-Align quality, AES, and DeQA gains for two endpoint-ascent steps over one; CLIP is higher for J = 1.Against the proxy, the two-step endpoint target improves Q-Align quality by 0.0125, Q-Align aesthetics by 0.0215, and DeQA by 0.0100.Source: §F.4 · p. 21
• 10.3% lower update time comes from one endpoint-target ascent step rather than two.Mean update times are 414.1 seconds for endpoint J = 1 and 461.7 seconds for endpoint J = 2.Source: §F.4 · p. 21

G Additional Experimental Results

• Figure S1 compares the base model and three trained methods across four scenes, illustrating differences in material texture, lighting, and local detail.The scenes are a museum interior, linen jacket, bicycle, and reading room.Source: G Additional Experimental Results · p. 21, Figure S1 · p. 22
• VOG increases runtime by 1.46×–1.98× and peak GPU memory by 1.21×–1.72× relative to CFG, with per-device sampling fitting on one NVIDIA RTX 4090 GPU.Table S9 reports runtime and peak memory relative to CFG.Source: G Additional Experimental Results · p. 22, Table S9 · p. 23
• FID falls from 7.58 to 4.44 on DiT-XL/2 and from 6.20 to 3.89 on SiT-XL/2 at 25 steps; the relative DiT-XL/2 reduction grows from 19% at 50 steps to 41% at 25 steps.The architectures trade Precision against Recall in opposite directions at 25 steps.Source: G Additional Experimental Results · pp. 22–23, Table S10 · p. 23
• VOG improves F1-scores and lowers Wasserstein distances in both branching-tree settings, while mitigating high-scale coverage loss and improving FID when composed with guidance heuristics.The reported examples include Recall ≈0.58 where CFG falls below 0.50, and Interval CFG FID falling from 1.637 to 1.523 on EDM2-L.Source: G Additional Experimental Results · p. 23, Table S11 · p. 24, Figure S3 · p. 24

H Reward and Evaluation Histories

• Reward-history plots separate sampled-endpoint rewards from late-state proxies; raw update means and trailing means are shown separately, with the update-20 reference marked.The training reward combines HPSv3++ and CLIP; ReFL and DRTune score late-state predictions, while DVOC and DiffusionOPSD score sampled endpoints.Source: H Reward and Evaluation Histories · p. 25, Figure S5 · p. 26

H.1 Online Evaluation

• DVOC’s online weighted evaluation reward rises from 0.766 at update 20 to 0.911 at update 100, while ImageReward rises from 1.267 to 1.366.ImageReward is held out of training; the online evaluation condition and noise sampling differ from the paired main-table evaluation.Source: §H.1 · p. 26, Figure S6 · p. 27

H.2 Active-Step Histories

• Larger K improves later weighted evaluation reward in the reported runs, while K = 1 retains a higher late-training AES score; ImageReward has no consistent checkpoint ordering.The K variants have matching initialization scores on all four displayed evaluators.Source: §H.2 · pp. 26, 28, Figure S7 · p. 27