Unifying Deep Stochastic Processes for Image Enhancement

Abstract: Deep stochastic processes have recently become a central paradigm for image enhancement, with many methods explicitly conditioning the stochastic trajectory on the degraded input. However, the relationship between these conditional processes and standard diffusion models remains unclear. In this work, we introduce a unified perspective on stochastic image enhancement by classifying recent methods into three families of continuous-time processes: unconditional diffusion models, Ornstein-Uhlenbeck (OU) processes, and diffusion bridges. We show that all of these approaches arise from a common stochastic differential equation (SDE) formulation. This framework makes explicit that seemingly disparate methods differ primarily in their drift and diffusion terms, terminal distributions, and boundary conditions, while schedulers and samplers constitute orthogonal design choices. Leveraging this unification, we conduct a controlled empirical study across multiple image enhancement tasks using identical architectures and training protocols. Our results reveal no consistently dominant method; instead, we identify and disentangle the specific design choices that most strongly influence performance. Finally, we release ItoVision, a modular PyTorch library that implements the unified framework and enables rapid prototyping and fair comparison of stochastic image enhancement methods.

• The paper unifies unconditional diffusion, Ornstein–Uhlenbeck processes, diffusion bridges, and flow matching within one continuous-time SDE framework, showing that BBDM, DDBM-VE, and I²SB are equivalent processes under different formulations.
• The paper’s controlled evaluation across five image-enhancement tasks finds no universally dominant method, with performance depending more on temperature, sampling stochasticity, endpoint regularization, and architecture than on process family.
• The paper identifies low-temperature collinearity in InDI-like processes and oversmoothing from deterministic bridge samplers, recommending ancestral sampling and careful temperature selection for reliable restoration quality.

Overview

The paper "Unifying Deep Stochastic Processes for Image Enhancement" (2605.01568) addresses a persistent methodological problem in diffusion-based image restoration: the proliferation of conditional stochastic processes—ResShift, InDI, IR-SDE, BBDM, DDBM, I$^2$SB, GOUB, UniDB—whose reported advantages are confounded with scheduler, sampler, discretization, and architecture choices. The authors' central contribution is a continuous-time SDE framework that expresses all of these methods, plus unconditional diffusion and flow matching, as instances of a single forward SDE

$$\mathrm{d}\mathbf{x}_t = f(\mathbf{x}_t, t, \mathbf{y})\,\mathrm{d}t + g(t)\,\mathrm{d}\mathbf{w}_t,$$

differing only in drift, diffusion coefficient, terminal distribution, and boundary conditions. Under this formulation, schedulers, samplers, and discretizations become orthogonal design choices rather than defining characteristics of a method. The paper pairs this unification with a large-scale controlled empirical study across five enhancement tasks using identical backbones, losses ($x_0$-prediction with MSE), and training protocols, and releases the resulting implementation as the modular PyTorch library ItoVision.

A unified taxonomy of stochastic processes

The framework organizes methods into three families. Unconditional processes (DM-VE/VP, flow matching) do not involve the degraded input $y$ in the forward process; conditioning is achieved only by feeding $y$ to the score network. Ornstein–Uhlenbeck processes (IR-SDE, ResShift, InDI) use mean-reverting drift $\beta_t(\mathbf{y} - \mathbf{x}_t)$ toward a Gaussian terminal distribution $\mathcal{N}(y, \tau^2 I)$ centered on the degraded image. Diffusion bridges (BBDM, DDBM-VE/VP, I$^2$SB, GOUB, UniDB) enforce arrival at $y$ at $t=1$ almost surely, via Doob's $$\mathrm{d}\mathbf{x}_t = f(\mathbf{x}_t, t, \mathbf{y})\,\mathrm{d}t + g(t)\,\mathrm{d}\mathbf{w}_t,$$0-transform or, in UniDB's case, a stochastic optimal control objective whose terminal cost recovers the $$\mathrm{d}\mathbf{x}_t = f(\mathbf{x}_t, t, \mathbf{y})\,\mathrm{d}t + g(t)\,\mathrm{d}\mathbf{w}_t,$$1-transform as $$\mathrm{d}\mathbf{x}_t = f(\mathbf{x}_t, t, \mathbf{y})\,\mathrm{d}t + g(t)\,\mathrm{d}\mathbf{w}_t,$$2.

Several propositions establish equivalences that dissolve apparent distinctions between methods:

• Flow matching as an SDE: with scheduler $$\mathrm{d}\mathbf{x}_t = f(\mathbf{x}_t, t, \mathbf{y})\,\mathrm{d}t + g(t)\,\mathrm{d}\mathbf{w}_t,$$3 and diffusion coefficient $$\mathrm{d}\mathbf{x}_t = f(\mathbf{x}_t, t, \mathbf{y})\,\mathrm{d}t + g(t)\,\mathrm{d}\mathbf{w}_t,$$4, the forward SDE has marginals $$\mathrm{d}\mathbf{x}_t = f(\mathbf{x}_t, t, \mathbf{y})\,\mathrm{d}t + g(t)\,\mathrm{d}\mathbf{w}_t,$$5, recovering rectified-flow transitions exactly.
• ResShift in continuous time: an OU process with $$\mathrm{d}\mathbf{x}_t = f(\mathbf{x}_t, t, \mathbf{y})\,\mathrm{d}t + g(t)\,\mathrm{d}\mathbf{w}_t,$$6 induces transition densities identical to ResShift's marginalized Markov chain, showing that ResShift and IR-SDE differ only in their diffusion coefficients.
• InDI as generalized flow matching: an OU process with $$\mathrm{d}\mathbf{x}_t = f(\mathbf{x}_t, t, \mathbf{y})\,\mathrm{d}t + g(t)\,\mathrm{d}\mathbf{w}_t,$$7 reproduces InDI's marginals; InDI is to flow matching what IR-SDE is to standard diffusion.
• BBDM = DDBM-VE = I$$\mathrm{d}\mathbf{x}_t = f(\mathbf{x}_t, t, \mathbf{y})\,\mathrm{d}t + g(t)\,\mathrm{d}\mathbf{w}_t,$$8SB: applying the $$\mathrm{d}\mathbf{x}_t = f(\mathbf{x}_t, t, \mathbf{y})\,\mathrm{d}t + g(t)\,\mathrm{d}\mathbf{w}_t,$$9-transform to a Wiener process (equivalently VE diffusion with constant scheduler) yields the Brownian bridge; deriving I$x_0$0SB's bridge from Schrödinger bridge theory under paired-data assumptions yields the same SDE. Three methods previously regarded as fundamentally incomparable are shown to be the same process differing only in scheduler choice.

The paper also derives closed-form solutions for the cosine scheduler integral, previously computed numerically in IR-SDE, GOUB, and UniDB, enabling arbitrary-timestep inference without parameter recomputation.

Controlled empirical study

Experiments cover facial super-resolution (FFHQ, ×8), general super-resolution (DIV2K, ×4), low-light enhancement (LOL), colorization (ImageNet, latent diffusion), and deraining (Rain1400), with UNet backbones (31M–119M parameters) shared across methods within each task, plus a DiT-backbone ablation. All models use MSE loss with $x_0$1 prediction and ancestral sampling at NFE = 35 for the main comparison.

The headline finding is negative and challenges a common assumption in the literature: no consistently dominant method emerges, and plain diffusion matches or outperforms specialized conditional processes on several tasks. Aggregated z-scores place ResShift first and DM-VP second; in facial super-resolution and low-light enhancement, unconditional diffusion achieves the best results outright. Family-level analysis shows task dependence: unconditional processes favor super-resolution and low-light enhancement, bridges favor colorization, and both are competitive for deraining. This mirrors earlier findings in GANs and diffusion design-space studies where training and architectural choices dominate method-specific differences.

Two concrete failure mechanisms are identified and disentangled:

• Low-temperature collinearity: InDI's default temperature ($x_0$2) causes intermediate states to concentrate near the segment between $x_0$3 and $x_0$4. The network then learns to extrapolate the displacement $x_0$5 as a function of time instead of modeling the true conditional score, propagating its own predictions during inference. Critically, lowering ResShift's temperature from its default $x_0$6 to InDI's level degrades ResShift comparably, demonstrating that the failure is a property of the temperature regime, not of any particular method. Unconditional processes are immune because $x_0$7 is independent of $x_0$8, so extrapolation bias cannot form.
• Deterministic sampling causes oversmoothing in bridges: since bridge trajectories start at fixed points and Brownian motion is their only stochasticity source, deterministic samplers (Euler ODE, exponential integrator, mean-reverting ODE, second-order Runge–Kutta) collapse trajectories onto straight lines converging to the conditional mean. No deterministic sampler achieves LPIPS below 0.3 on super-resolution even at NFE = 100, whereas ancestral sampling remains the most consistent choice globally. Discretization strategy, by contrast, has no significant effect.

Additional analyses show UniDB's endpoint regularization (finite $x_0$9) explains its gains over GOUB, which shares the same process; the DiT backbone generally underperforms UNet; and running each method with its original configuration mostly yields worse results than the unified setup—for instance, DDPM improves when switching from $y$0-prediction to $y$1-prediction—reinforcing that implementation choices, not process definitions, drive much of the reported variance.

Limitations and open questions

The paper's scope excludes several relevant classes: guidance-based restoration methods that require knowledge of the degradation operator, Schrödinger bridge methods designed for unpaired data, and architecture-focused improvements such as Refusion. The equivalence results rely on paired training data, which holds for the benchmarks studied but not for all translation settings. The empirical conclusion that conditional trajectories offer no inherent advantage is established under one training protocol (MSE, $y$2-prediction); whether alternative losses or parameterizations could change the ranking is not resolved, although the original-setup experiments suggest some sampler–parameterization combinations interact favorably. The temperature analysis identifies low $y$3 as harmful but does not derive an optimal temperature selection rule. Finally, the finding that performance degrades with increasing NFE under the DiT backbone is observed but left unexplained.

Conclusion

This work reduces a fragmented landscape of conditional stochastic image enhancement methods to a single SDE formalism, proves several previously unrecognized equivalences—notably that BBDM, DDBM-VE, and I$y$4SB are the same process—and demonstrates through controlled experimentation that the underlying process definition matters far less than temperature, sampling stochasticity, and endpoint regularization. The identification of low-temperature collinearity and deterministic-sampler oversmoothing as concrete, mechanistic failure modes provides actionable guidance, and the released ItoVision library makes fair comparison reproducible going forward.