Efficient Image Restoration with State-Dependent Forward Diffusion

TMLR 2026

Ziwei Luo, Fredrik K. Gustafsson, Jens Sjölund, Thomas B. Schön
  Uppsala University     AstraZeneca  
Paper arXiv Code

State-Dependent Foorward Diffusion (FoD)


FoD introduces the mean reversion term $\mu - x_t$ into both the drift and diffusion functions, enabling high-quality data samples with a single diffusion process.

Summary

This paper proposes to perform image restoration through a state-dependent mean-reverting \textbf{fo}rward \textbf{d}iffusion (FoD) process. In contrast to traditional diffusion-based approaches that rely on a coupled forward-backward diffusion scheme, FoD directly learns image restoration through a single forward diffusion process, yielding a simple yet efficient framework. The core of FoD is a state-dependent stochastic differential equation (SDE) that involves a mean-reverting term in both the drift and diffusion functions. This mean-reverting structure drives the low-quality data toward the clean endpoint with controlled stochastic variation, therefore simulating a stochastic interpolation between source and target distributions. More importantly, FoD is analytically tractable and is trained using a simple stochastic flow matching objective, enabling few-step sampling during inference. The proposed FoD model, despite its simplicity, achieves strong overall performance on various image restoration tasks compared to representative diffusion, diffusion bridge, and flow matching approaches.

State-Dependent Forward Diffusion Process

$\mathrm{d}x_t = \theta_t \, (\mu - x_t) \mathrm{d}t + \sigma_t (x_t - \mu) \mathrm{d}w_t$


where $\mu \sim p_\text{data}$ is the clean data and $x_0 \sim p_\text{prior}$ is the source data. The diffusion volatility increases in the beginning steps and then decreases to zero when $x_t$ converges to $\mu$.


Stochastic Flow Matching

The FoD process is analytically tractable and follows a multiplicative stochastic structure. We show that this model can be learned by approximating the vector field from each noisy state to the final clean data, called the Stochastic Flow Matching:


$L_\text{SFM}(\phi)=\mathbb{E}_{\mu,x_t} [ \| (\mu - x_t) - f_\phi(x_t, t) \|^2 ]$

Algorithms

The standard training and sampling (via the Euler–Maruyama method) procedures are provided in Algorithm 1 and Algorithm 2. In addition, we also provide fast sampling with Markov and non-Markov chains in Algorithm 3 and Algorithm 4.

Connection to Prior Work

Stochastic Interpolants: Let us recall the solution of the FoD process. By setting the initial state to $x_0$ and rearranging the equation, we obtain a stochastic process in the interpolant form:


$x_t = I(t, x_0, \mu) = x_0 \, \alpha_t + \mu \ (1 - \alpha_t), \quad \alpha_t = e^{-\int_0^t \bigl(\theta_z + \frac{1}{2}\sigma_z^2 \bigr) \mathrm{d}{z} + \int_0^t{\sigma}_z {\mathrm{d}{w}_z}}.$


Note: FoD can be viewed as a powerful instantiation of SI, distinguished by two key properties: multiplicative log-normal interpolation and a state-dependent stochastic path from $x_0$ to $\mu$.


Flow Matching: We consider a deterministic version of the FoD process, i.e., omitting the diffusion term or setting $\sigma_t = 0$ for all times. This gives a mean-reverting ODE that bridges two distributions without noise injection:


$\mathrm{d}{x}_t = \theta_t \, (\mu - x_t) \mathrm{d}{t}, \quad x_t = \bigl(x_s - \mu \bigr) \, e^{-\int_{s}^t \theta_z \mathrm{d}{z}} + \mu$


Note: Our primary FoD model can be also regarded as a stochastic extension of flow matching models.

Results

Image Restoration



Image-to-Image Translation


Fast Sampling with Markov and non-Markov Chains

Ablation Experiments on Noise Injection

Thanks for your interest!

BibTeX

If our code helps your research or work, please consider citing our paper. The following are BibTeX references:


      @article{
      luo2026efficient,
      title={Efficient Image Restoration with State-Dependent Forward Diffusion},
      author={Ziwei Luo and Fredrik K. Gustafsson and Jens Sj{\"o}lund and Thomas B. Sch{\"o}n},
      journal={Transactions on Machine Learning Research},
      issn={2835-8856},
      year={2026},
      url={https://openreview.net/forum?id=Eq9k6Va3hY},
      note={}
      }