Stochastic Optimal Control Sampling for Diffusion Inverse Problems

Conference: ECCV 2026
Paper: ECCV Official
Code: https://github.com/zjqwq01/SOCS-DIP
Area: Image Restoration
Keywords: Diffusion Models, Inverse Problems, Stochastic Optimal Control, Langevin Sampling, Flow Matching

TL;DR

By formulating the reverse denoising process as a controlled dynamical system, this paper derives a closed-form single-step stochastic optimal control update with bounded terminal penalties, bypassing full-trajectory second-order Hessian computation while adaptively balancing data fidelity and generative prior preservation.

Background & Motivation

Unsupervised image inverse problem solving using pretrained diffusion models has emerged as a prominent paradigm in computational imaging and image restoration. Forward image degradations are often many-to-one and highly nonlinear, making the reconstruction of clean images from degraded measurements an ill-posed inverse problem with non-unique solutions. Pretrained diffusion models capture rich high-order structural priors and natural image manifold geometry, enabling high-quality posterior sampling by conditioning reverse denoising dynamics on degraded measurements without task-specific retraining.

However, existing training-free guided diffusion frameworks face a dilemma. Mainstream Bayesian posterior sampling schemes, such as DPS and DAPS, heuristically incorporate the gradient of a quadratic data-fidelity term to steer sampling toward measurements, which implicitly presumes an infinite terminal penalty. Such aggressive likelihood injection frequently pulls intermediate states off the underlying generative prior manifold during early-to-mid stages, leading to visual artifacts and oversmoothed textures. Conversely, previous attempts to apply stochastic optimal control (SOC) to diffusion inverse problems, such as DOC, formulate the problem via classical iterative Linear Quadratic Regulator (iLQR) trajectory optimization; optimizing over the entire diffusion path requires second-order Taylor updates and Hessian computations, resulting in prohibitive GPU memory overhead and slow sampling latencies that prevent application to modern high-resolution models.

The fundamental tension stems from the fact that an uncontrolled diffusion model already produces high-quality samples on its native manifold; inverse guidance should inject minimal-energy, bounded control to steer trajectories toward measurement-consistent sets rather than overwhelming the generative prior. Hence, a lightweight sampling scheme with rigorous optimal control guarantees that decouples into explicit single-step updates is urgently needed. The core idea is to frame reverse-diffusion dynamics as a stochastic optimal control problem with bounded terminal measurement penalties, derive a closed-form analytical transport update via the certainty equivalence principle, and utilize Langevin energy sampling with closed-form pullbacks at each step to project measurement-consistent estimates back onto the denoising flow without trajectory-level Hessians.

Method

Overall Architecture

SOCS models the reverse denoising process as a linear stochastic differential equation (SDE) system. By injecting a control signal \(u_{t,\gamma}\) into the drift, the objective steers the terminal state \(x_T\) to minimize the Euclidean observation residual \(\frac{\gamma}{2}\|y - \mathcal{H}(x_T)\|_2^2\) while simultaneously penalizing the control energy \(\int_0^T \frac{1}{2}\|u_{t,\gamma}\|_2^2 dt\). At each discrete outer denoising step \(t_i\), the pipeline first advances the probability-flow ordinary differential equation (ODE) using a few-step Euler solver to obtain a coarse clean image estimate \(x_T|t_i\). Depending on the conditioning and linearity of degradation operator \(\mathcal{H}\), the estimate is aligned with measurements via either closed-form affine subspace projection or bounded Langevin energy sampling. Finally, the refined estimate is analytically mapped back onto the reverse denoising flow at \(x_{t_{i+1}}\) via the derived closed-form transport equation.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["Degraded measurement y and initial noise x_t0"] --> B["Stage lookahead: few-step Euler ODE solver<br/>Fast forward extrapolation yields coarse estimate x_T|ti"]
    B --> C{"Degradation operator decoupling and dual-mode sampling<br/>Evaluate operator local linearity and conditioning"}
    C -->|Locally linear and well-conditioned| D["Closed-form affine subspace projection<br/>Direct mapping onto measurement-consistent set"]
    C -->|Complex nonlinear or severely ill-posed| E["Bounded terminal penalty modulation and energy sampling<br/>Langevin dynamics data-consistency energy refinement"]
    D --> F["Single-step SOC closed-form pullback<br/>Analytical state transport maps estimate back onto manifold"]
    E --> F
    F -->|Iterate discrete reverse steps to termination| G["Reconstructed image aligned with measurements x_N"]

Key Designs

1. Degradation operator decoupling and dual-mode sampling: bridging linear projection and modern generative paradigms Addressing the limitation that conventional sampling schemes struggle to handle diverse degradation operators under a unified pipeline, SOCS introduces an operator-aware branched sampling strategy and extends it to latent diffusion models and flow-matching backbones. Under well-conditioned linear degradation models (e.g., super-resolution), the Jacobian satisfies \(J_\mathcal{H} \approx \mathcal{H}\); taking \(\gamma \to \infty\) eliminates iterative inner refinement by projecting directly via pseudoinverse: $\(\hat{x}_T(x_t) = (I - \mathcal{H}^\dagger \mathcal{H}) x_T(x_t) + \mathcal{H}^\dagger(y)\)$ This projection preserves measurement consistency along observed dimensions while leaving null-space textures to the generative prior. For latent diffusion models (SD v1.5) and flow-matching frameworks (SD3-medium), rectified flow trajectories approximate linear ODEs with near-constant velocity under straight-line interpretations, allowing SOCS to be deployed plug-and-play in high-resolution latent spaces without modifying pretrained weights.

2. Bounded terminal penalty modulation and energy sampling: preserving manifold geometry under explicit measurement constraints Standard Bayesian guidance typically assumes \(\gamma \to \infty\), imposing infinite terminal penalties that distort denoising trajectories in ill-posed or nonlinear regimes (such as phase retrieval and nonlinear deblurring). SOCS explicitly retains a finite penalty coefficient \(\gamma\), recasting reverse control as an energy-bounded minimization problem. In nonlinear regimes, the regularized Gauss-Newton normal matrix inverse is computed via implicit Langevin energy minimization: $\(\hat{x}_T^{j+1} = \hat{x}_T^j - \eta \nabla_{\hat{x}_T^j} \left\| \left(I + \gamma(1 - e^{2\bar{f}_T}) J_\mathcal{H}(x_T^u)^\top\right) \circ \mathcal{H}(\hat{x}_T^j) - J_\mathcal{H}(x_T^u)^\top (y - e^{\bar{f}_T} \mathcal{H}(x_0)) \right\|_2^2 + \sqrt{2\eta}\epsilon_j\)$ Iterative stochastic gradient refinement circumvents explicit matrix inversion. Incorporating a finite \(\gamma\) establishes an adaptive regularization mechanism that drives convergence toward observed measurements while leaving structural degrees of freedom for diffusion prior dynamics.

3. Single-step SOC closed-form pullback: bypassing trajectory Hessians via analytical transport dynamics Prior control-theoretic approaches for inverse problems maintain full-trajectory rollouts using iterative LQR, requiring second-order Taylor expansions and Hessian computations that cause GPU memory bottlenecks. SOCS proves that under linear SDE dynamics \(dx_t = (f_t x_t + g_t u_{t,\gamma})dt + g_t dw_t\) and quadratic terminal costs, the certainty equivalence principle guarantees that the optimal controller matches that of the deterministic ODE system. The authors derive closed-form analytical expressions for the optimal control \(u_{t,\gamma}^*\) and state trajectory \(x_t\): $\(x_t = e^{\bar{f}_t} x_0 + \gamma e^{\bar{f}_t} e^{\bar{f}_T} \bar{g}_t^2 d_{0,\gamma} J_\mathcal{H}(x_T^u)^\top \left(y - e^{\bar{f}_T}\mathcal{H}(x_0)\right)\)$ where \(d_{0,\gamma} = (I + \gamma e^{2\bar{f}_T}\bar{g}_T^2 J_\mathcal{H}^\top \circ \mathcal{H})^{-1}\). This closed-form transport operator pulls the measurement-consistent clean estimate \(\hat{x}_T\) back onto the reverse SDE denoising manifold at each discrete step. Operating strictly with first-order Jacobian-vector products, SOCS eliminates trajectory-wide rollouts and Hessian storage, reducing computational latency to the level of standard sampling-based baselines.

A Worked Example

Consider \(4\times\) bicubic super-resolution (linear) and phase retrieval (nonlinear) under additive Gaussian noise: 1. Initialization: Given degraded observation \(y\) with noise level \(\beta_y = 0.05\), sample initial state \(x_{t_0} \sim \mathcal{N}(0, I)\). 2. State lookahead prediction: At outer step \(i\), the score network advances \(x_{t_i}\) across 4 Euler ODE steps to compute coarse estimate \(x_T|t_i\). 3. Consistency refinement: - Super-resolution: In SOCS-Linear mode, the estimate is directly projected onto the affine subspace via \(\hat{x}_T|t_i = (I - \mathcal{H}^\dagger \mathcal{H}) x_T|t_i + \mathcal{H}^\dagger y\). - Phase Retrieval: Because magnitude-only measurements lack a bounded linear pseudoinverse, execution branches to SOCS-Nonlinear-γ with \(\gamma = 10^7\), performing \(N_{\text{inner}}\) Langevin gradient refinement steps to obtain \(\hat{x}_T|t_i\). 4. Single-step manifold pullback: Substituting refined \(\hat{x}_T|t_i\) into the analytical state transition formula calculates the next reverse state \(x_{t_{i+1}}\). 5. Iteration & convergence: Repeating this decouple-refine-pullback sequence across 250 to 500 outer steps yields clean reconstructions aligned with measurements.

Key Experimental Results

Main Results

The method was evaluated on 100 validation images each from FFHQ 256×256 and ImageNet 256×256 across 6 linear and 3 nonlinear tasks. Metrics include PSNR, SSIM, LPIPS, and FID. Evaluations were also conducted at 768×768 resolution using SD3-medium.

Dataset & Task	Metrics	SOCS-Linear (Ours)	SOCS-Nonlinear-γ (Ours)	DAPS (Baseline)	DPS (Baseline)	DDRM / DDNM (Classical)
FFHQ: Super-Resolution 4×	PSNR / SSIM / FID	28.51 / 0.797 / 72.99	28.66 / 0.800 / 76.86	28.31 / 0.768 / 74.75	24.25 / 0.677 / 84.33	27.53 / 0.796 / 92.19 (DDRM)
FFHQ: Random Inpainting	PSNR / SSIM / FID	29.31 / 0.834 / 77.09	30.63 / 0.872 / 52.73	29.66 / 0.800 / 54.57	28.30 / 0.812 / 71.74	29.17 / 0.850 / 64.65 (DDNM)
FFHQ: Box Inpainting	PSNR / SSIM / FID	24.22 / 0.794 / 53.21	24.45 / 0.826 / 43.79	24.41 / 0.753 / 51.51	23.49 / 0.804 / 62.18	24.14 / 0.837 / 48.96 (DDNM)
ImageNet: Super-Resolution 4×	PSNR / SSIM / FID	26.42 / 0.733 / 93.77	26.36 / 0.708 / 97.54	25.95 / 0.667 / 112.60	22.56 / 0.574 / 155.01	25.02 / 0.694 / 158.35 (DDNM)
ImageNet: Gaussian Deblur	PSNR / SSIM / FID	25.38 / 0.720 / 111.49	26.67 / 0.740 / 110.69	26.37 / 0.697 / 109.61	22.84 / 0.587 / 115.02	27.81 / 0.639 / 124.55 (DDRM)
FFHQ: Nonlinear Deblur	PSNR / SSIM / FID	21.73 / 0.647 / 113.90	28.42 / 0.786 / 68.20	28.10 / 0.762 / 69.45	23.05 / 0.637 / 90.69	29.16 / 0.783 / 63.22 (RED-diff)
FFHQ: Phase Retrieval	PSNR / SSIM / FID	N/A (No linear pseudoinverse)	29.59 / 0.808 / 44.86	29.88 / 0.807 / 45.33	16.29 / 0.442 / 137.81	16.83 / 0.521 / 209.75 (RED-diff)

Ablation Study

The table below details comparisons on SD3-medium at 768×768 resolution for super-resolution ×12, along with sensitivity to the terminal penalty parameter \(\gamma\).

Model & Config	Task	PSNR (↑)	SSIM (↑)	LPIPS (↓)	FID (↓)	Note
FlowDPS	SD3 ×12 (Bicubic)	29.73	0.775	0.482	61.24	Vector field guidance with limited perceptual fidelity
FlowChef	SD3 ×12 (Bicubic)	29.56	0.846	0.346	152.00	Severe prior deviation degrades perceptual FID
LatentDAPS	SD3 ×12 (Bicubic)	29.65	0.819	0.276	65.41	Decoupled annealing lacking optimal transport pullback
SOCS-SD3 (Ours)	SD3 ×12 (Bicubic)	30.45	0.842	0.248	45.72	Optimal control guidance achieves superior fidelity and perception
SOCS-Nonlinear (\(\gamma \to \infty\))	Inverse problems general	Slight drop	Slight drop	Higher	Higher	Overfitting measurement noise leads to oversmoothing
SOCS-Nonlinear-γ (\(\gamma = 10^7\))	Inverse problems general	Best	Best	Lowest	Lowest	Saturated regularized balance between prior and data fidelity
SOCS-Nonlinear-γ (\(\gamma = 10^3\))	Inverse problems general	Major drop	Major drop	Substantially higher	Major increase	Insufficient guidance strength fails to enforce consistency

Key Findings

• Robust performance plateau across control parameters: Varying \(\gamma\) from \(10^5\) to \(10^9\) yields a stable performance plateau; \(\gamma \le 10^3\) under-injects measurement guidance, while \(\gamma \to \infty\) drives trajectories off the prior manifold, validating the bounded penalty design.
• Operator conditioning governs variant selection: For well-conditioned linear problems, SOCS-Linear provides computational speed and stability via affine subspace projections; for ill-posed, large null-space, or nonlinear settings, SOCS-Nonlinear-\(\gamma\) with Jacobian-aware Langevin updates substantially outperforms linear baselines.
• Direct optimization path with lower deviation budget: Owing to closed-form derivations, SOCS avoids iterative trajectory rollouts and achieves lower data-consistency error than DAPS under equivalent cumulative deviation budgets.

Highlights & Insights

• Control-theoretic foundation for empirical guidance gradients: The paper demonstrates that the empirical gradient guidance term \(\nabla_{x_t} \|H(x_t) - y\|_2^2\) commonly used in Bayesian diffusion inverse sampling corresponds to the infinite penalty limit (\(\gamma \to \infty\)) of stochastic optimal control.
• Trajectory iLQR reduced to closed-form single-step transport: Applying certainty equivalence allows the optimal control policy to be solved analytically, eliminating full-trajectory rollouts and Hessian storage and matching the efficiency of first-order sampling methods.
• Seamless generalization to rectified flow matching: The SOC formulation extends from linear SDE diffusions to flow-matching architectures (SD3-medium), validating plug-and-play restoration capability across different generative frameworks.

Limitations & Future Work

• Empirical tuning of inner Langevin steps: For highly nonlinear operators, inner Langevin refinement still requires multiple iterations, increasing overall neural function evaluations (NFEs).
• Reliance on local Gauss-Newton linearization: In complex nonlinear regimes, early-stage Gauss-Newton approximations can deviate from the true manifold under high noise levels; adaptive curvature corrections could be explored.
• Scope bounded by Gaussian noise models: The current derivation explicitly presumes additive Gaussian measurement noise and known degradation models; extending control formulations to blind operators and Poisson noise remains future work.

Related Work & Insights

• vs DPS / DAPS: DPS relies on heuristic Gaussian likelihood approximations; DAPS decouples annealing without optimal control energy constraints. SOCS derives analytical state transport under minimal control energy, preventing manifold drift via bounded \(\gamma\).
• vs DOC / iLQR-based control: DOC performs iterative trajectory rollouts requiring second-order Hessian computation; SOCS derives single-step closed-form updates, lowering compute and memory costs to first-order Jacobian complexity.
• vs DDRM / DDNM: DDRM and DDNM are restricted to linear inverse problems with accessible pseudoinverses; SOCS provides unified formulations spanning linear projections and nonlinear problems like phase retrieval.

Rating

• Novelty: ⭐⭐⭐⭐⭐ Establishes a rigorous analytical single-step closed-form SOC framework for diffusion inverse sampling.
• Experimental Thoroughness: ⭐⭐⭐⭐⭐ Extensively verified across VE/VP-SDE, LDMs, and SD3 Flow-Matching over 9 linear and nonlinear benchmarks.
• Writing Quality: ⭐⭐⭐⭐⭐ Rigorous mathematical derivation paired with structured and clear explanations.
• Value: ⭐⭐⭐⭐⭐ Establishes an efficient, theoretically grounded baseline for training-free diffusion inverse problems.