Forward-only Diffusion Probabilistic Models

## Abstract

Abstract This work presents a forward-only diffusion (FoD) approach for generative modelling. In contrast to traditional diffusion models that rely on a coupled forward-backward diffusion scheme, FoD directly learns data generation through a single forward diffusion process, yielding a simple yet efficient generative framework. The core of FoD is a state-dependent linear stochastic differential equation that involves a mean-reverting term in both the drift and diffusion functions. This mean-reversion property guarantees the convergence to clean data, naturally 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 a few-step non-Markov chain sampling during inference. The proposed FoD model, despite its simplicity, achieves competitive performances on various image-conditioned (e.g., image restoration) and unconditional generation tasks, demonstrating its effectiveness in generative modelling. Our code is available at https://github.com/Algolzw/FoD .

## 1 Introduction

The diffusion model has become a central theme in generative modelling (Sohl-Dickstein et al., 2015; Ho et al., 2020; Song et al., 2021; Karras et al., 2022). A crucial feature of the diffusion model is the use of a forward process that gradually perturbs the data into noise, coupled with a backward process that learns to transform noise back to data Sohl-Dickstein et al. (2015); Ho et al. (2020). Benefiting from this forward-backward framework, the diffusion models have achieved remarkable performance in producing high-quality results across a wide range of applications, including image synthesis (Dhariwal and Nichol, 2021; Saharia et al., 2022b; Rombach et al., 2022; Peebles and Xie, 2023; Podell et al., 2024), translation (Meng et al., 2022; Saharia et al., 2022a; Tumanyan et al., 2023; Kawar et al., 2023; Brooks et al., 2023), and restoration Saharia et al. (2022c); Kawar et al. (2022); Luo et al. (2023a); Wang et al. (2024); Lin et al. (2024).

Despite this success, the reliance on the coupled forward-backward construction substantially increases the algorithmic complexity at the same time as it leads to a challenging model training problem. In addition, the necessity of corrupting data to noise in diffusion models further imposes an undesirable constraint for image-conditioned generation (Kawar et al., 2022; Saharia et al., 2022a), where ideally the generative process should start with image conditions that are structurally more informative than noise (Luo et al., 2023a; Saharia et al., 2022c; Liu et al., 2023). This naturally leads to a fundamental question:

“Could a simpler, single diffusion process suffice for effective generative modelling?”

This paper answers this question affirmatively by introducing a probabilistic forward-only diffusion model (FoD). Our exploration starts from the mean-reverting stochastic differential equation (SDE) Gillespie (1996); Luo et al. (2023a), where the data is stochastically driven toward a specified state characterized by a fixed mean and variance. Notably, setting the mean to zero recovers the standard forward diffusion process Song et al. (2021). Inspired by this, we propose a new form of the mean-reverting SDE, which adds mean-reversion to both the drift and diffusion functions as a state-dependent diffusion process. Here, we highlight the mean-reversion diffusion function as it guarantees the convergence to the noise-free mean state. By setting the mean to the target data, FoD naturally simulates the data transition between source and target distributions, without requiring a separate backward process.

This paper further demonstrates that the FoD process is analytically tractable and follows a multiplicative stochastic structure. Moreover, we show that the model can be learned by approximating the vector field from each noisy state to the final clean data, a processes we refer to as stochastic flow matching. The result is a simple, yet effective training process. Based on the tractable solution and the flow matching objective, FoD enables a few-step sampling strategy with both Markov and non-Markov chains, enabling more efficient data generation without compromising sample quality.

In addition, as a closely related work of our method, it is worth noting that flow matching (Lipman et al., 2022; Liu et al., 2022) can also eliminate the need for a separate data perturbation process by modelling a continuous flow from source distribution to target distribution. However, the noise injection, which has been shown crucial in generative models Song and Ermon (2019), is also eliminated due to the modelling of ordinary differential equations (ODEs). As a result, its performance drops significantly when handling image-conditioned generation tasks, such as image restoration (IR), which aims to recover high-quality images from their degraded low-quality counterparts (Albergo et al., 2023a; Martin et al., 2024; Ohayon et al., 2024). In contrast, our FoD is naturally a stochastic extension of flow matching, which avoids the issue mentioned above by simulating SDEs with a state-dependent diffusion process, making it well-suited for image-conditioned generation. Notably, omitting the diffusion term of FoD yields an ODE variant of our model, which is a special case of flow matching and can be applied for unconditional generation.

In experiments, we evaluate the proposed FoD model on various IR tasks, including image deraining, dehazing, low-light enhancement, and inpainting, as well as an unconditional generation task on CIFAR-10. The results show that FoD achieves competitive performance across all tasks and datasets, demonstrating its effectiveness in generative modelling.

## 2 Background

Given a source distribution $p_{\text{prior}}$ and an unknown target data distribution $p_{\text{data}}$, our goal is to build a probability
path $\{p(x_{t})\}_{t=0}^{T}$ that transports between the source distribution $p(x_{0})=p_{\text{prior}}$ and the target distribution $p(x_{T})=p_{\text{data}}$. In this paper, the source can be either noise, for unconditional generation, or images, for image-conditioned generation, e.g., image restoration.

### 2.1 Diffusion Models

Given a target data point $x_{T}\sim p_{\text{data}}$, diffusion models Sohl-Dickstein et al. (2015); Ho et al. (2020) define a Markov chain forward process to progressively perturb the data into noise ($x_{T}\rightarrow x_{0}$) and then learn its reverse process to reconstruct the data ($x_{0}\rightarrow x_{T}$). This coupled forward-backward process can be defined by stochastic differential equations (SDEs) Song et al. (2021), given by:

$$ $\underbrace{\mathop{}\!\mathrm{d}x_{t}=f(x_{t},t)\mathop{}\!\mathrm{d}t+g(t)\mathop{}\!\mathrm{d}w_{t}}_{\text{Forward process}}\quad\text{and}\quad\underbrace{\mathop{}\!\mathrm{d}x_{t}=\bigl{[}f(x_{t},t)-g(t)^{2}\,\nabla\log p_{t}(x_{t})\bigr{]}\mathop{}\!\mathrm{d}t+g(t)\mathop{}\!\mathrm{d}\bar{w}_{t}}_{\text{Backward process}},$ (1) $$

where $f(x,t)$ is the drift function and $g(t)$ is the diffusion function. Furthermore, $w$ and $\bar{w}$ are the standard Wiener process and its reverse process, respectively. We use $p_{t}(x_{t})$ to denote the marginal probability density of $x_{t}$. The term $\nabla\log p_{t}(x_{t})$, called the score function, is the sought-after objective in the backward (also called the reverse-time) SDE, which is often learned by a time-dependent neural network Song et al. (2021) via score-matching. The training objective can also be converted to learn noise matching as in DDPMs Ho et al. (2020). Moreover, the source distribution in diffusion models is often a Gaussian with a predefined mean and variance. In order to generate high-quality samples, diffusion models typically require hundreds or even thousands of sampling steps

### 2.2 Flow Matching Generative Models

Flow matching (Lipman et al., 2022; Albergo and Vanden-Eijnden, 2022; Liu et al., 2022) is a simple regression objective used for learning the velocity field $v(x_{t},t)$ that transports a sample $x_{t}$ from the source distribution to the target distribution along the probability path $p(x_{t})$ Lipman et al. (2024).
More specifically, flow matching models aim to learn the ordinary differential equation (ODE): $\mathop{}\!\mathrm{d}{x_{t}}=v(x_{t},t)\mathop{}\!\mathrm{d}{t}$, where $x_{0}\sim p_{\text{prior}}$ and
the drift $v(x_{t},t)$ transports samples from $x_{0}$ to $x_{1}\sim p_{\text{data}}$. Here, each latent variable $x_{t}$ in the ODE path is drawn by linearly interpolating source and target data samples, i.e., $x_{t}=tx_{1}+(1-t)x_{0}$. Then the training can be performed by uniformly sampling data pairs and timesteps and optimizing a flow matching objective, as:

$$ $L_{\text{FM}}(\phi)=\mathbb{E}_{x_{0},x_{1},t\sim\mathcal{U}(0,1)}\bigl{[}\|(x_{1}-x_{0})-v_{\phi}(x_{t},t)\|^{2}\bigr{]},$ (2) $$

where $v_{\phi}(x_{t},t)$ is a neural network approximating the true velocity field.
Flow matching models eliminate the diffusion term from the generative process and thus lead to a simpler and more direct learning procedure based on ODE paths. However, in this paper, we observe that applying it to image-conditioned generation tasks, such as image restoration, leads to a significant performance drop due to the lack of stochastic noise injection (see Section [4.1](#S4.SS1) for more details).

## 3 Forward-only Diffusion Probabilistic Models

An overview of the probabilistic forward-only diffusion (FoD) model is shown in Figure [1](#S3.F1). In the following, we first introduce the mean-reverting SDE in Section [3.1](#S3.SS1). Motivated by the mean-reverting property of this SDE, we propose FoD and provide its training objective in Section [3.2](#S3.SS2) and Section [3.3](#S3.SS3), respectively. Finally, in Section [3.4](#S3.SS4), we analyze the ODE variant of our model and show how to learn it for unconditional generation.

[Uncaptioned image]: /html/2505.16733/assets/x1.png

### 3.1 Preliminaries: Mean-reverting SDE

Our exploration starts from a mean-reverting SDE Gillespie (1996); Luo et al. (2023a) where the data is stochastically driven towards a state characterized by a specified mean $\mu$ and variance $\lambda^{2}$, defined by

$$ $\mathop{}\!\mathrm{d}{x}_{t}=\theta_{t}\,(\mu-x_{t})\mathop{}\!\mathrm{d}{t}+\sigma_{t}\mathop{}\!\mathrm{d}{w}_{t},$ (3) $$

where $\{\theta_{t}\}_{t=0}^{T}$ and $\{\sigma_{t}\}_{t=0}^{T}$ are positive mean-reversion and diffusion schedules, respectively. By coupling the schedules as $\sigma_{t}^{2}\,/\,\theta_{t}=2\,\lambda^{2}$ for all timesteps, the following solution is obtained Luo et al. (2023a)

$$ $x_{t}=\mu+\bigl{(}x_{0}-\mu\bigr{)}\,\mathrm{e}^{-\int^{t}_{0}\theta_{z}\mathop{}\!\mathrm{d}z}+\int^{t}_{0}\sigma_{z}\,\mathrm{e}^{-\int^{t}_{s}\theta_{s}\mathop{}\!\mathrm{d}s}\mathop{}\!\mathrm{d}w_{z}.$ (4) $$

As $t\to\infty$, the SDE converges to a stationary state $x_{T}\sim\mathcal{N}(x_{t}\mid\mu,\lambda^{2})$. This property suggests constructing a process that transports samples from the source distribution $p_{\text{prior}}$ to the target distribution $p_{\text{data}}$, by setting the mean $\mu$ to be a sample from $p_{\text{data}}$. However, as can be observed from Eq. ([4](#S3.E4)), the resulting sample $x_{T}$ is still noisy, with variance $\lambda^{2}$, which works against our goal of generating high-quality clean data samples. In the following sections, we address this problem by introducing mean-reversion in both the drift and diffusion functions.

### 3.2 Forward-only Diffusion Process

We begin by designing an SDE with mean-reversion terms in both the drift and diffusion functions, as

$$ $\mathop{}\!\mathrm{d}{x}_{t}=\theta_{t}\,(\mu-x_{t})\mathop{}\!\mathrm{d}{t}+\sigma_{t}\,(x_{t}-\mu)\mathop{}\!\mathrm{d}{w}_{t}.$ (5) $$

This is a state-dependent linear SDE with multiplicative noise, where the diffusion volatility increases in the beginning steps and then decreases to zero when $x_{t}$ converges to $\mu$. We typically use $x_{t}-\mu$ rather than $\mu-x_{t}$ in the diffusion function, such that the SDE simulates a reverse Wiener process as in diffusion models Song et al. (2021). For image generation, the noise $\mathop{}\!\mathrm{d}w_{t}$ is added independently at each pixel, meaning that this SDE is applied for image transitions pixel-by-pixel, under the Itô interpretation.

We refer to this SDE as the forward-only diffusion (FoD) process, and present its solution as follows:

###### Proposition 3.1 .

Given an initial state $x_{s}$ at time $s<t$, the unique solution to the SDE ([5](#S3.E5)) is

$$ $x_{t}=\bigl{(}x_{s}-\mu\bigr{)}\,\mathrm{e}^{-\int_{s}^{t}\bigl{(}\theta_{z}+\frac{1}{2}\sigma_{z}^{2}\bigr{)}\mathop{}\!\mathrm{d}{z}+\int_{s}^{t}{\sigma}_{z}{\mathop{}\!\mathrm{d}{w}_{z}}}+\mu,$ (6) $$

where the stochastic integral is interpreted in the Itô sense and can be reparameterised as $\bar{\sigma}_{s:t}\,\epsilon$, where $\bar{\sigma}_{s:t}=\sqrt{\int_{s}^{t}\sigma_{z}^{2}\mathop{}\!\mathrm{d}z}$, and $\epsilon\sim\mathcal{N}(0,I)$ is a standard Gaussian noise.

The proof is provided in Appendix [A.1](#A1.SS1). In addition, the solution in Eq. ([6](#S3.E6)) shows that the stochastic flow field $\mu-x_{t}$ forms a Geometric Brownian motion Ross (2014) and yields the following corollary:

###### Corollary 3.2 .

Under the same assumptions as in Proposition [3.1](#S3.Thmtheorem1), the stochastic flow field $\mu-x_{t}$ satisfies the multiplicative stochastic structure. More precisely, it is log-normally distributed by

$$ $\log(\mu-x_{t})\sim\mathcal{N}\Bigl{(}\log(\mu-x_{s})-\int_{s}^{t}\bigl{(}\theta_{z}+\frac{1}{2}\sigma_{z}^{2}\bigr{)}\mathop{}\!\mathrm{d}z,\,\int_{s}^{t}\sigma_{z}^{2}\,\mathop{}\!\mathrm{d}z\,I\Bigr{)}.$ (7) $$

This follows directly from Proposition [3.1](#S3.Thmtheorem1), by rearranging $\mu-x_{t}$ to the left of Eq. ([6](#S3.E6)) and applying the logarithm to both sides (see Appendix [A.1](#A1.SS1)). The subtractive form of the logarithm reflects that the flow field decays multiplicatively from its initial value with a stochastic exponential scaling.

Notational clarifications: Although the sign of $\mu-x_{t}$ in general can be either positive or negative, it remains consistent across all times $t$ for a given sample; therefore, we choose to omit absolute values inside the logarithmic terms in Eq. ([7](#S3.E7)) for notational convenience. In addition, we further let $\bar{m}_{s:t}=-\int_{s}^{t}(\theta_{z}+\frac{1}{2}\sigma_{z}^{2})\mathop{}\!\mathrm{d}z$ and $\bar{m}_{t}=\bar{m}_{0:t}$ in the rest of the paper to simplify the notation.

###### Proposition 3.1 .

###### Corollary 3.2 .

### 3.3 Stochastic Flow Matching

Let us now explain how we can learn this FoD process, i.e., transforming data from a known source distribution $p_{\text{prior}}$ to an unknown target distribution $p_{\text{data}}$. Inspired by the diffusion model Ho et al. (2020), we define our FoD model as $p_{\phi}(x_{0:T})$, a joint distribution with learnable transitions starting at $x_{0}$, as

$$ $p_{\phi}(x_{0:T})=p_{\text{prior}}(x_{0})\prod_{t=0}^{T-1}p_{\phi}(x_{t+1}\mid x_{t}),\qquad x_{0}\sim p_{\text{prior}}.$ (8) $$

We propose to set the transition kernel $p_{\phi}(x_{t+1}|x_{t})$ to be in the same log-Gaussian form as Eq. ([6](#S3.E6)). The training can then be performed by minimizing the negative log-likelihood of $p_{\phi}(x_{T})$, which is equivalent to optimizing the following objective:

$$ $\mathbb{E}_{p}\Bigl{[}\,\sum_{t=0}^{T-1}D_{KL}(p(x_{t+1}\mid x_{t},x_{T})\,\|\,p_{\phi}(x_{t+1}\mid x_{t}))\,\Bigr{]}.$ (9) $$

The proof is provided in Appendix [A.2](#A1.SS2). During training, we set $x_{T}$ equal to $\mu$ such that the SDE ([5](#S3.E5)) converges to data $\mu$ exactly. Then, the conditional distribution $p(x_{t+1}|x_{t},x_{T}=\mu)$ is tractable as shown in Eq. ([6](#S3.E6)). By letting the functions $f_{\mu}(x_{t})=\mu-x_{t}$ and $f_{\phi}(x_{t},t)=\hat{\mu}_{\phi}-x_{t}$ denote the ground truth and the model prediction of the stochastic flow field, respectively, we transform the distributions in Eq. ([9](#S3.E9)) from SDE states to stochastic flow fields. Note that this transformation, i.e., from $p(\cdot|x_{t},\mu)$ to $p(\cdot|f_{\mu}(x_{t}))$, holds because its Jacobian determinant equals one. Instead of Eq. ([9](#S3.E9)), we can therefore minimize the KL divergence between two stochastic flow distributions:

$$ $\mathbb{E}_{p}\Bigl{[}\,\sum_{t=0}^{T-1}D_{KL}(p(f_{\mu}(x_{t+1})\mid f_{\mu}(x_{t})\,\|\,p(f_{\phi}(x_{t+1},t)\mid f_{\phi}(x_{t},t))\,\Bigr{]}.$ (10) $$

Combining this with Corollary [3.2](#S3.Thmtheorem2), we obtain the final objective:

$$ $\displaystyle L_{\text{SFM}}(\phi)$ $\displaystyle\coloneqq\mathbb{E}_{\mu\sim p_{\text{data}},x_{t}\sim q(x_{t}|x_{0},\mu)}\Bigl{[}\,\|\log(\mu-x_{t})-\log f_{\phi}(x_{t},t)\|^{2}\,\Bigr{]}$ (11) $\displaystyle\ \approx\mathbb{E}_{\mu\sim p_{\text{data}},x_{t}\sim q(x_{t}|x_{0},\mu)}\Bigl{[}\|(\mu-x_{t})-f_{\phi}(x_{t},t)\|^{2}\Bigr{]},$ $$

where the approximation follows from a first-order Taylor expansion close to the optimum.
Please refer to Appendix [A.3](#A1.SS3) for more details.
This objective is referred to as stochastic flow matching and it is entirely linear, which leads to a simple and numerically stable training process.

The standard training and sampling (via the Euler–Maruyama method) procedures are provided in Algorithm [1](#alg1) and Algorithm [2](#alg2). In addition, the target data estimate $\hat{\mu}$ is given by

$$ $\hat{\mu}=x_{t}+f_{\phi}(x_{t},t),$ (12) $$

which can be applied to the forward transition ([6](#S3.E6)) for fast data sampling.

#### Fast Sampling with Markov and non-Markov Chains

While the generation can be performed by iteratively solving the SDE ([5](#S3.E5)) with numerical schemes such as the Euler–Maruyama method, it usually requires hundreds of sampling steps. Fortunately, the tractable solution of FoD naturally enables fast sampling during inference, by choosing times discretely with a larger step size $k$, as $t=[0,k,2k,3k,\dots,T]$, where $T$ is the total number of timesteps. As our prediction at each step is the sought-after target data $\hat{\mu}$, as in Eq. ([12](#S3.E12)), the next state $x_{t+k}$ can be sampled following Eq. ([6](#S3.E6)) with either Markov or non-Markov chains, by setting the transition to $x_{t}\rightarrow x_{t+k}$ or $x_{0}\rightarrow x_{t+k}$, as illustrated in Algorithm [3](#alg3) and Algorithm [4](#alg4), respectively. More discussions are provided in Section [5.1](#S5.SS1).

Figure: Algorithm 1 FoD Training

### 3.4 Learning Mean-reverting ODEs for Unconditional Generation

Although FoD works well for image-conditioned generation, the multiplicative structure of the forward process also poses a challenge for unconditional generation, where the source distribution is often a Gaussian. Then, mixing the source Gaussian noise with our injected log-Gaussian noise (from Eq. ([6](#S3.E6))) significantly complicates the learning and worsens the sample quality. To address this, we consider a deterministic version of FoD ([5](#S3.E5)), omitting the diffusion term or setting $\sigma_{t}=0$ for all times. As a result, we obtain a mean-reverting ODE that transforms source data to target data without the multiplicative noise injection, as

$$ $\mathop{}\!\mathrm{d}{x}_{t}=\theta_{t}\,(\mu-x_{t})\mathop{}\!\mathrm{d}{t}\quad\text{with solution}\quad x_{t}=\bigl{(}x_{s}-\mu\bigr{)}\,\mathrm{e}^{-\int_{s}^{t}\theta_{z}\mathop{}\!\mathrm{d}{z}}+\mu.$ (13) $$

By setting $s=0$ and rearranging the solution, we obtain an interpolation between source and target:

$$ $x_{t}=\alpha_{t}\ x_{0}+(1-\alpha_{t})\ x_{T},$ (14) $$

where $x_{T}=\mu$ and $\alpha_{t}=\mathrm{e}^{-\int_{0}^{t}\theta_{z}\mathop{}\!\mathrm{d}{z}}$ is a time-related coefficient decreasing from 1 to 0. This equation follows a similar interpolation as in flow matching Lipman et al. (2022) but with a special transportation path given by $\theta_{t}\,(\mu-x_{t})$. We can then directly learn the drift, resulting in a conditional flow matching objective:

$$ $L_{\text{CFM}}\coloneqq\mathbb{E}_{\mu,x_{t}}\Bigl{[}\|(\mu-x_{t})-f_{\phi}(x_{t},t)\|^{2}\Bigr{]}=\mathbb{E}_{\mu,x_{t}}\Bigl{[}\|\alpha_{t}(\mu-x_{0})-f_{\phi}(x_{t},t)\|^{2}\Bigr{]},$ (15) $$

which is a deterministic form of stochastic flow matching ([11](#S3.E11)).
In practice, we can directly learn the straight path $\mu-x_{0}$ and define the $\alpha$ schedule to be linear or even a constant, e.g. $\alpha_{t}=1$ for all time $t$, in which case our mean-reverting ODE naturally becomes rectified flow Liu et al. (2022) exactly. In other words, our primary FoD model can be regarded as a stochastic extension of flow matching models.

## 4 Experiments

Our experiments mainly focus on the image restoration (IR) task, which is a typical image-conditioned generation problem aiming to recover high-quality images from their degraded low-quality counterparts. The results on various datasets are reported in Section [4.1](#S4.SS1). In addition, we provide an experiment on CIFAR-10 Krizhevsky et al. (2009) in Section [4.2](#S4.SS2) to illustrate unconditional generation(^1^11We provide more implementation details, datasets, and additional results in the appendix).

#### Setup

We use a U-Net Ronneberger et al. (2015) architecture similar to DDPM Ho et al. (2020) for flow prediction in all tasks. Attention layers are removed for efficient training and testing, similar to IR-SDE Luo et al. (2023a, b). We choose the commonly used cosine and linear schedules Nichol and Dhariwal (2021) for $\theta_{t}$ and $\sigma_{t}$, respectively, and normalize $\sigma_{t}^{2}$ to sum to 1 to ensure numerical stability under multiplicative noise perturbation. The number of sampling steps is fixed to 100 for all tasks. We use the AdamW Loshchilov and Hutter (2017) optimizer with parameters $\beta_{1}=0.9$ and $\beta_{2}=0.99$. The training requires 500 thousand iterations with a learning rate of $10^{-4}$. All our models are trained on an A100 GPU with 40 GB of memory for about 1.5 days.

For image restoration, the Learned Perceptual Image Patch Similarity (LPIPS) Zhang et al. (2018) and Fréchet inception distance (FID) Heusel et al. (2017) are reported to evaluate the perceptual fidelity and visual quality. Additionally, Peak Signal-to-Noise Ratio (PSNR) and Structural Similarity Index Measure (SSIM) Wang et al. (2004) are also included to measure the pixel/structure similarity. For unconditional image generation, FID is adopted as the primary metric for assessing sample quality.

Figure: Table 1: Quantitative comparison between our method with other diffusion and flow matching approaches on four different image restoration datasets. Here, ‘FoD w/ MC’ and ‘FoD w/ NMC’ denote 10-step Markov and non-Markov chain fast sampling, respectively. ‘FoD w/o noise’ represents the ODE variant of our model.

Figure: Figure 2: Comparison of FoD with other approaches on four IR tasks. We also provide the results for ‘FoD w/o noise’ in the last column to illustrate the importance of noise injection in image restoration.
Refer to caption: /html/2505.16733/assets/x2.png

### 4.1 Image Restoration

We evaluate our method on four image restoration tasks: 1) image deraining on the Rain100H dataset (Yang et al., 2017), 2) dehazing on the RESIDE-6k dataset (Qin et al., 2020), 3) low-light enhancement on LOL (Wei et al., 2018), and 4) face inpainting on CelebA-HQ Karras et al. (2017).

#### Comparison Approaches

In all experiments, we select IR-SDE Luo et al. (2023a) as the main comparison method to evaluate the performance gap between forward-backward and forward-only schemes for diffusion-based image restoration. We also compare with other flow-based approaches such as Rectified flow Liu et al. (2022); Liu (2022) and posterior-mean rectified flow (PMRF) Ohayon et al. (2024) that learn ODEs and also allow the model to generate images with a single forward process. In addition, a U-Net model, using the same architecture as our FoD, is trained with the $\ell_{1}$ loss as a CNN baseline on all tasks for reference.

#### Results

The quantitative comparisons on four IR tasks are reported in Table [1](#S4.T1). The proposed FoD achieves the best results across all datasets in comparison to other diffusion-based and flow-based approaches.
Compared to the U-Net baseline, IR-SDE successfully improves the results on perceptual metrics (LPIPS and FID) across all tasks, proving the effectiveness of the forward-backward based diffusion IR schemes. We observe that flow-based approaches, such as ReFlow and PMRF, perform inferiorly on all distortion and perceptual metrics.
While PMRF improves the PSNR results for flow matching-based IR, the performance gain potentially comes from the two-stage training strategy and the small noise injection in the initial state of rectified flow. Moreover, Table [1](#S4.T1) presents additional results for methods ‘FoD w/o noise’, ‘FoD w/ MC’, and ‘FoD w/ NMC’, where the former represents the ODE variant of our model (Section [3.4](#S3.SS4)) and the latter two methods mean we adopt 10-step Markov and non-Markov chain sampling strategies in inference (see Algorithm [3](#alg3) and Algorithm [4](#alg4)). More results and comparisons for these FoD variants will be discussed in Section [5](#S5).

We also provide visual comparisons in Figure [2](#S4.F2), showing that our FoD produces the most realistic and high-fidelity
results. In particular, while all deterministic approaches without noise injection (ReFlow, PMRF and ‘FoD w/o noise’) tend to generate over-smooth results (see e.g. the left eye area in the face inpainting case), the proposed FoD model consistently generates sharp images.

### 4.2 Unconditional Image Generation

We evaluate the unconditional generation performance of our model on the CIFAR-10 dataset Krizhevsky et al. (2009). Specifically, we showcase results from both our FoD model and its ODE-based variant, FoD-ODE (see Section [3.4](#S3.SS4)). Both models share the same $\theta$ schedule and are sampled with 100 steps. We compare against several baselines, including 1) diffusion models with forward-backward schemes such as DDPM Ho et al. (2020) and Score SDE Song et al. (2021), as well as 2) forward-only frameworks: score-based generative models (SGMs) like NSCN Song and Ermon (2019) and NSCNv2 Song and Ermon (2020), flow matching generative models using diffusion and optimal transport (OT) paths Lipman et al. (2022), and Rectified Flow Liu et al. (2022) which also adopts the OT path for data transformation.

Figure: Figure 3: Left: Visual results of unconditional generation on the CIFAR-10 dataset by two variants of our FoD model. Right: Quantitative comparison of our methods with other approaches on CIFAR-10.
Refer to caption: /html/2505.16733/assets/x3.png

#### Results

We present the generated image samples and quantitative comparisons in Figure [3](#S4.F3), with visual examples on the left and numerical results summarized in the table on the right. Under the forward-only framework, our FoD model with SDE sampling achieves a competitive FID of 7.89, outperforming other score-based generative models as well as the flow matching approach along the diffusion path. The ODE-based variant of FoD further improves the performance to an FID of 5.01, surpassing the flow matching model based on the optimal transport path. Nonetheless, our overall performance on CIFAR-10 remains inferior to Rectified Flow and conventional diffusion models using a forward-backward scheme. We partly attribute this performance gap to differences in architectural choices, hyperparameter tuning, etc. However, we also note that the main aim of this paper is to construct an effective single diffusion process model, which we expect to be particularly well-suited for image-conditioned generation tasks such as image restoration.

## 5 Discussion

In this section, we analyze the performance change of using a non-Markov chain for fast sampling and illustrate the effectiveness of using noise injection in image restoration. Then, we provide examples of our FoD sampling to further illustrate the data transformation process in different generation tasks.

### 5.1 Fast Sampling with Markov and Non-Markov Chains

Figure: Figure 4: Comparison of fast sampling with Markov chain (MC) vs. with non-Markov chain (NMC) on deraining task, using the same pretrained model.
Refer to caption: /html/2505.16733/assets/x4.png

As discussed in Section [3.3](#S3.SS3), our method naturally supports fast sampling using both Markov and non-Markov chains. Table [1](#S4.T1) demonstrates that with only 10 sampling steps, the non-Markov variant (FoD w/ NMC) achieves even better results than the standard Euler–Maruyama solver, without compromising perceptual quality. Note that the key difference between faster sampling approaches and the standard FoD sampling is that the latter’s iterative process highly depends on the $\mu$ prediction, while the former gradually refines the state $x_{t}$ rather than $\mu$. To further investigate the impact of sampling steps, we provide a detailed analysis in Figure [4](#S5.F4). For comparison, the Euler–Maruyama solver with 100 steps is included as a baseline. We observe that both sampling strategies yield improved distortion metrics as the number of steps decreases, especially in the low-step areas (5–20), where they substantially outperform the baseline with a small perceptual performance drop. Similar trends are observed for LPIPS and FID, but one can observe that the results on these metrics slightly degraded when reducing the number of steps from 20 to 5. More comparisons are provided in Appendix [C.3](#A3.SS3). The overall results suggest that using 10 sampling steps can serve as a practical rule of thumb for fast sampling with both Markov and non-Markov chains.

### 5.2 Effectiveness of Noise Injection

In this section, we emphasize the importance of noise injection in image-conditioned generation. For this purpose, we provide an additional experiment to train a noise-free FoD on four image restoration tasks, by simply letting $\sigma_{t}=0$ for all times. Then our model turns out to be an ODE model as described in Section [3.4](#S3.SS4). We keep the $\theta$ schedule the same as FoD for a fair comparison. The quantitative results are shown in Table [1](#S4.T1), as ‘FoD w/o noise’. As can be observed, the results drop significantly across all tasks and metrics. The corresponding visual results are also provided in the last column of Figure [2](#S4.F2), where the produced images are blurry and unclear compared to those of our original FoD. Moreover, the training curves of FoD with and without noise injection are provided in Figure [5](#S5.F5). One can observe that FoD consistently outperforms its noise-free variant on all tasks, further demonstrating the effectiveness and importance of noise injection in image-conditioned generation.

Figure: Figure 5: Training curves of our FoD and its noise-free variant on different tasks in terms of PSNR and LPIPS. All these results demonstrate the effectiveness of noise injection in image restoration.
Refer to caption: /html/2505.16733/assets/x5.png

### 5.3 Illustration of the Diffusion Process

To clearly show the noise injection process with our model, we apply the trained FoD on various tasks, including both image restoration and unconditional image generation, and illustrate their intermediate states along timesteps as shown in Figure [6](#S5.F6). We find that the noise is only injected into the degraded areas, such as the masked regions in deraining and face inpainting. This means our model focuses more on restoring the degradations instead of reconstructing the whole image, yielding a more efficient solution to the image restoration problem. More examples for the FoD diffusion process are provided in Appendix [C.3](#A3.SS3).

[Uncaptioned image]: /html/2505.16733/assets/x13.png

## 6 Related Work

Denoising diffusion models (Sohl-Dickstein et al., 2015; Ho et al., 2020; Song et al., 2021; Karras et al., 2022) and flow matching models (Lipman et al., 2022; Liu et al., 2022; Lipman et al., 2024; Albergo et al., 2023a; Gat et al., 2024) are two popular frameworks in generative modelling and have been widely applied to various applications including image generation (Dhariwal and Nichol, 2021; Ho et al., 2022; Peebles and Xie, 2023; Ho and Salimans, 2022; Ma et al., 2024), text-to-image generation (Rombach et al., 2022; Ruiz et al., 2023; Podell et al., 2024; Saharia et al., 2022b), image translation (Meng et al., 2022; Lugmayr et al., 2022; Saharia et al., 2022a; Su et al., 2022; Liu et al., 2023; Ben-Hamu et al., 2024), etc. Inspired by their success in producing photo-realistic images conforming to human preference, these models have recently been applied to image restoration for advanced performance (Wang et al., 2024; Saharia et al., 2022c; Yue et al., 2023; Luo et al., 2023a; Kawar et al., 2022; Luo et al., 2023b; Yue et al., 2024; Luo et al., 2024a, b). In addition, Albergo et al. (2023a) unify diffusion and flow matching models through stochastic interpolants, a general framework for constructing stochastic processes that interpolate between arbitrary distributions. This formulation has also been applied to image restoration Albergo et al. (2023b), where stochastic flows guided by corrupted observations recover clean images, effectively serving as a stochastic extension of flow matching for inverse problems. Subsequent works enhance this approach using pretrained flow matching models Ben-Hamu et al. (2024) or refined training pipelines Ohayon et al. (2024). These works are closely related to ours, but they all adopt noise-free generation processes. In contrast, we build FoD with a state-dependent diffusion process for image generation, which is naturally well-suited for image-conditioned generation, and which reduces to deterministic flow matching when the diffusion term vanishes.

## 7 Conclusion

This paper introduces a new framework, named FoD, for generative modelling with a single forward diffusion process. We have shown that FoD is analytically tractable and can be trained using a simple flow matching objective. FoD is evaluated on various image-conditioned generation tasks as well as an unconditional generation task, achieving a competitive performance compared to other diffusion-based and flow-based approaches, demonstrating its effectiveness and efficiency in generative modelling, particularly in the application of image restoration.

## Acknowledgements

This research was supported by the Wallenberg AI, Autonomous Systems and Software Program (WASP) funded by the Knut and Alice Wallenberg Foundation;
by the project Deep Probabilistic Regression – New Models and Learning Algorithms (contract number: 2021-04301) funded by the Swedish Research Council; and by the
Kjell & Märta Beijer Foundation.
The computations were enabled by the Berzelius resource provided by the Knut and Alice Wallenberg Foundation at the National Supercomputer Centre.

## References

- Albergo and Vanden-Eijnden [2022]
Michael S Albergo and Eric Vanden-Eijnden.
Building normalizing flows with stochastic interpolants.
*arXiv preprint arXiv:2209.15571*, 2022.
- Albergo et al. [2023a]
Michael S Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden.
Stochastic interpolants: A unifying framework for flows and diffusions.
*arXiv preprint arXiv:2303.08797*, 2023a.
- Albergo et al. [2023b]
Michael S Albergo, Mark Goldstein, Nicholas M Boffi, Rajesh Ranganath, and Eric Vanden-Eijnden.
Stochastic interpolants with data-dependent couplings.
*arXiv preprint arXiv:2310.03725*, 2023b.
- Ben-Hamu et al. [2024]
Heli Ben-Hamu, Omri Puny, Itai Gat, Brian Karrer, Uriel Singer, and Yaron Lipman.
D-flow: Differentiating through flows for controlled generation.
*arXiv preprint arXiv:2402.14017*, 2024.
- Brooks et al. [2023]
Tim Brooks, Aleksander Holynski, and Alexei A Efros.
Instructpix2pix: Learning to follow image editing instructions.
In *Proceedings of the IEEE/CVF conference on computer vision and pattern recognition*, pages 18392–18402, 2023.
- Dhariwal and Nichol [2021]
Prafulla Dhariwal and Alexander Nichol.
Diffusion models beat gans on image synthesis.
*Advances in neural information processing systems*, 34:8780–8794, 2021.
- Gat et al. [2024]
Itai Gat, Tal Remez, Neta Shaul, Felix Kreuk, Ricky TQ Chen, Gabriel Synnaeve, Yossi Adi, and Yaron Lipman.
Discrete flow matching.
*Advances in Neural Information Processing Systems*, 37:133345–133385, 2024.
- Gillespie [1996]
Daniel T Gillespie.
Exact numerical simulation of the ornstein-uhlenbeck process and its integral.
*Physical review E*, 54(2):2084, 1996.
- Heusel et al. [2017]
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter.
GANs trained by a two time-scale update rule converge to a local nash equilibrium.
In *Proceedings of Advances in Neural Information Processing Systems (NeurIPS)*, volume 30, 2017.
- Ho and Salimans [2022]
Jonathan Ho and Tim Salimans.
Classifier-free diffusion guidance.
*arXiv preprint arXiv:2207.12598*, 2022.
- Ho et al. [2020]
Jonathan Ho, Ajay Jain, and Pieter Abbeel.
Denoising diffusion probabilistic models.
*Advances in neural information processing systems*, 33:6840–6851, 2020.
- Ho et al. [2022]
Jonathan Ho, Chitwan Saharia, William Chan, David J Fleet, Mohammad Norouzi, and Tim Salimans.
Cascaded diffusion models for high fidelity image generation.
*Journal of Machine Learning Research*, 23(47):1–33, 2022.
- Karras et al. [2017]
Tero Karras, Timo Aila, Samuli Laine, and Jaakko Lehtinen.
Progressive growing of gans for improved quality, stability, and variation.
*arXiv preprint arXiv:1710.10196*, 2017.
- Karras et al. [2022]
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine.
Elucidating the design space of diffusion-based generative models.
*Advances in neural information processing systems*, 35:26565–26577, 2022.
- Kawar et al. [2022]
Bahjat Kawar, Michael Elad, Stefano Ermon, and Jiaming Song.
Denoising diffusion restoration models.
*Advances in Neural Information Processing Systems*, 35:23593–23606, 2022.
- Kawar et al. [2023]
Bahjat Kawar, Shiran Zada, Oran Lang, Omer Tov, Huiwen Chang, Tali Dekel, Inbar Mosseri, and Michal Irani.
Imagic: Text-based real image editing with diffusion models.
In *Proceedings of the IEEE/CVF conference on computer vision and pattern recognition*, pages 6007–6017, 2023.
- Krizhevsky et al. [2009]
Alex Krizhevsky, Geoffrey Hinton, et al.
Learning multiple layers of features from tiny images.
Technical report, Toronto, ON, Canada, 2009.
- Lin et al. [2024]
Xinqi Lin, Jingwen He, Ziyan Chen, Zhaoyang Lyu, Bo Dai, Fanghua Yu, Yu Qiao, Wanli Ouyang, and Chao Dong.
Diffbir: Toward blind image restoration with generative diffusion prior.
In *European Conference on Computer Vision*, pages 430–448. Springer, 2024.
- Lipman et al. [2022]
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le.
Flow matching for generative modeling.
*arXiv preprint arXiv:2210.02747*, 2022.
- Lipman et al. [2024]
Yaron Lipman, Marton Havasi, Peter Holderrieth, Neta Shaul, Matt Le, Brian Karrer, Ricky TQ Chen, David Lopez-Paz, Heli Ben-Hamu, and Itai Gat.
Flow matching guide and code.
*arXiv preprint arXiv:2412.06264*, 2024.
- Liu et al. [2023]
Guan-Horng Liu, Arash Vahdat, De-An Huang, Evangelos Theodorou, Weili Nie, and Anima Anandkumar.
I ^2 sb: Image-to-image schrödinger bridge.
In *International Conference on Machine Learning*, pages 22042–22062. PMLR, 2023.
- Liu [2022]
Qiang Liu.
Rectified flow: A marginal preserving approach to optimal transport.
*arXiv preprint arXiv:2209.14577*, 2022.
- Liu et al. [2022]
Xingchao Liu, Chengyue Gong, and Qiang Liu.
Flow straight and fast: Learning to generate and transfer data with rectified flow.
*arXiv preprint arXiv:2209.03003*, 2022.
- Loshchilov and Hutter [2017]
Ilya Loshchilov and Frank Hutter.
Decoupled weight decay regularization.
*arXiv preprint arXiv:1711.05101*, 2017.
- Lugmayr et al. [2022]
Andreas Lugmayr, Martin Danelljan, Andres Romero, Fisher Yu, Radu Timofte, and Luc Van Gool.
Repaint: Inpainting using denoising diffusion probabilistic models.
In *Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition*, pages 11461–11471, 2022.
- Luo et al. [2023a]
Ziwei Luo, Fredrik K Gustafsson, Zheng Zhao, Jens Sjölund, and Thomas B Schön.
Image restoration with mean-reverting stochastic differential equations.
In *International Conference on Machine Learning*, pages 23045–23066. PMLR, 2023a.
- Luo et al. [2023b]
Ziwei Luo, Fredrik K Gustafsson, Zheng Zhao, Jens Sjölund, and Thomas B Schön.
Refusion: Enabling large-size realistic image restoration with latent-space diffusion models.
In *Proceedings of the IEEE/CVF conference on computer vision and pattern recognition*, pages 1680–1691, 2023b.
- Luo et al. [2024a]
Ziwei Luo, Fredrik K. Gustafsson, Zheng Zhao, Jens Sjölund, and Thomas B. Schön.
Controlling vision-language models for multi-task image restoration.
In *The Twelfth International Conference on Learning Representations*, 2024a.
- Luo et al. [2024b]
Ziwei Luo, Fredrik K Gustafsson, Zheng Zhao, Jens Sjölund, and Thomas B Schön.
Photo-realistic image restoration in the wild with controlled vision-language models.
In *Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition*, pages 6641–6651, 2024b.
- Ma et al. [2024]
Nanye Ma, Mark Goldstein, Michael S Albergo, Nicholas M Boffi, Eric Vanden-Eijnden, and Saining Xie.
Sit: Exploring flow and diffusion-based generative models with scalable interpolant transformers.
In *European Conference on Computer Vision*, pages 23–40. Springer, 2024.
- Martin et al. [2024]
Ségolène Martin, Anne Gagneux, Paul Hagemann, and Gabriele Steidl.
Pnp-flow: Plug-and-play image restoration with flow matching.
*arXiv preprint arXiv:2410.02423*, 2024.
- Meng et al. [2022]
Chenlin Meng, Yutong He, Yang Song, Jiaming Song, Jiajun Wu, Jun-Yan Zhu, and Stefano Ermon.
SDEdit: Guided image synthesis and editing with stochastic differential equations.
In *International Conference on Learning Representations*, 2022.
- Nichol and Dhariwal [2021]
Alexander Quinn Nichol and Prafulla Dhariwal.
Improved denoising diffusion probabilistic models.
In *International conference on machine learning*, pages 8162–8171. PMLR, 2021.
- Ohayon et al. [2024]
Guy Ohayon, Tomer Michaeli, and Michael Elad.
Posterior-mean rectified flow: Towards minimum mse photo-realistic image restoration.
*arXiv preprint arXiv:2410.00418*, 2024.
- Peebles and Xie [2023]
William Peebles and Saining Xie.
Scalable diffusion models with transformers.
In *Proceedings of the IEEE/CVF International Conference on Computer Vision*, pages 4195–4205, 2023.
- Podell et al. [2024]
Dustin Podell, Zion English, Kyle Lacey, Andreas Blattmann, Tim Dockhorn, Jonas Müller, Joe Penna, and Robin Rombach.
SDXL: Improving latent diffusion models for high-resolution image synthesis.
In *The Twelfth International Conference on Learning Representations*, 2024.
- Qin et al. [2020]
Xu Qin, Zhilin Wang, Yuanchao Bai, Xiaodong Xie, and Huizhu Jia.
FFA-Net: Feature fusion attention network for single image dehazing.
In *Proceedings of the AAAI Conference on Artificial Intelligence*, pages 11908–11915, 2020.
- Rombach et al. [2022]
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer.
High-resolution image synthesis with latent diffusion models.
In *Proceedings of the IEEE/CVF conference on computer vision and pattern recognition*, pages 10684–10695, 2022.
- Ronneberger et al. [2015]
Olaf Ronneberger, Philipp Fischer, and Thomas Brox.
U-net: Convolutional networks for biomedical image segmentation.
In *Medical image computing and computer-assisted intervention–MICCAI 2015: 18th international conference, Munich, Germany, October 5-9, 2015, proceedings, part III 18*, pages 234–241. Springer, 2015.
- Ross [2014]
Sheldon M Ross.
*Introduction to probability models*.
Academic press, 2014.
- Ruiz et al. [2023]
Nataniel Ruiz, Yuanzhen Li, Varun Jampani, Yael Pritch, Michael Rubinstein, and Kfir Aberman.
Dreambooth: Fine tuning text-to-image diffusion models for subject-driven generation.
In *Proceedings of the IEEE/CVF conference on computer vision and pattern recognition*, pages 22500–22510, 2023.
- Saharia et al. [2022a]
Chitwan Saharia, William Chan, Huiwen Chang, Chris Lee, Jonathan Ho, Tim Salimans, David Fleet, and Mohammad Norouzi.
Palette: Image-to-image diffusion models.
In *ACM SIGGRAPH 2022 conference proceedings*, pages 1–10, 2022a.
- Saharia et al. [2022b]
Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, et al.
Photorealistic text-to-image diffusion models with deep language understanding.
*Advances in neural information processing systems*, 35:36479–36494, 2022b.
- Saharia et al. [2022c]
Chitwan Saharia, Jonathan Ho, William Chan, Tim Salimans, David J Fleet, and Mohammad Norouzi.
Image super-resolution via iterative refinement.
*IEEE transactions on pattern analysis and machine intelligence*, 45(4):4713–4726, 2022c.
- Sohl-Dickstein et al. [2015]
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli.
Deep unsupervised learning using nonequilibrium thermodynamics.
In *International conference on machine learning*, pages 2256–2265. pmlr, 2015.
- Song and Ermon [2019]
Yang Song and Stefano Ermon.
Generative modeling by estimating gradients of the data distribution.
*Advances in neural information processing systems*, 32, 2019.
- Song and Ermon [2020]
Yang Song and Stefano Ermon.
Improved techniques for training score-based generative models.
*Advances in neural information processing systems*, 33:12438–12448, 2020.
- Song et al. [2021]
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole.
Score-based generative modeling through stochastic differential equations.
*International Conference on Learning Representations*, 2021.
- Su et al. [2022]
Xuan Su, Jiaming Song, Chenlin Meng, and Stefano Ermon.
Dual diffusion implicit bridges for image-to-image translation.
*arXiv preprint arXiv:2203.08382*, 2022.
- Tumanyan et al. [2023]
Narek Tumanyan, Michal Geyer, Shai Bagon, and Tali Dekel.
Plug-and-play diffusion features for text-driven image-to-image translation.
In *Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition*, pages 1921–1930, 2023.
- Wang et al. [2024]
Jianyi Wang, Zongsheng Yue, Shangchen Zhou, Kelvin CK Chan, and Chen Change Loy.
Exploiting diffusion prior for real-world image super-resolution.
*International Journal of Computer Vision*, 132(12):5929–5949, 2024.
- Wang et al. [2004]
Zhou Wang, Alan C Bovik, Hamid R Sheikh, and Eero P Simoncelli.
Image quality assessment: from error visibility to structural similarity.
*IEEE Transactions on Image Processing*, 13(4):600–612, 2004.
- Wei et al. [2018]
Chen Wei, Wenjing Wang, Wenhan Yang, and Jiaying Liu.
Deep retinex decomposition for low-light enhancement.
*arXiv preprint arXiv:1808.04560*, 2018.
- Yang et al. [2017]
Wenhan Yang, Robby T Tan, Jiashi Feng, Jiaying Liu, Zongming Guo, and Shuicheng Yan.
Deep joint rain detection and removal from a single image.
In *Proceedings of the IEEE conference on computer vision and pattern recognition*, pages 1357–1366, 2017.
- Yue et al. [2024]
Conghan Yue, Zhengwei Peng, Junlong Ma, Shiyan Du, Pengxu Wei, and Dongyu Zhang.
Image restoration through generalized ornstein-uhlenbeck bridge.
In *Proceedings of the 41st International Conference on Machine Learning*, pages 58068–58089, 2024.
- Yue et al. [2023]
Zongsheng Yue, Jianyi Wang, and Chen Change Loy.
Resshift: Efficient diffusion model for image super-resolution by residual shifting.
*Advances in Neural Information Processing Systems*, 36:13294–13307, 2023.
- Zhang et al. [2018]
Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang.
The unreasonable effectiveness of deep features as a perceptual metric.
In *Proceedings of IEEE Conference on Computer Vision and Pattern Recognition (CVPR)*, pages 586–595, 2018.

## Appendix A Proofs

### A.1 Proof to Proposition 3.1 and Corollary 3.2

###### Proposition A.1 .

Given an initial state $x_{s}$ at time $s<t$, the unique solution to the SDE ([5](#S3.E5)) is

$$ $x_{t}=\bigl{(}x_{s}-\mu\bigr{)}\,\mathrm{e}^{-\int_{s}^{t}\bigl{(}\theta_{z}+\frac{1}{2}\sigma_{z}^{2}\bigr{)}\mathop{}\!\mathrm{d}{z}+\int_{s}^{t}{\sigma}_{z}{\mathop{}\!\mathrm{d}{w}_{z}}}+\mu,$ (16) $$

where the stochastic integral is interpreted in the Itô sense and can be reparameterised as $\bar{\sigma}_{s:t}\,\epsilon$, where $\bar{\sigma}_{s:t}=\sqrt{\int_{s}^{t}\sigma_{z}^{2}\mathop{}\!\mathrm{d}z}$, and $\epsilon\sim\mathcal{N}(0,I)$ is a standard Gaussian noise.

###### Corollary A.2 .

Under the same assumptions as in Proposition [3.1](#S3.Thmtheorem1), the stochastic flow field $\mu-x_{t}$ satisfies the multiplicative stochastic structure. More precisely, it is log-normally distributed by

$$ $\log(\mu-x_{t})\sim\mathcal{N}\Bigl{(}\log(\mu-x_{s})-\int_{s}^{t}\bigl{(}\theta_{z}+\frac{1}{2}\sigma_{z}^{2}\bigr{)}\mathop{}\!\mathrm{d}z,\,\int_{s}^{t}\sigma_{z}^{2}\,\mathop{}\!\mathrm{d}z\,I\Bigr{)}.$ (17) $$

###### Proof.

Recall the FoD process from Eq. ([5](#S3.E5)):

$$ $\mathop{}\!\mathrm{d}{x}_{t}=\theta_{t}\,(\mu-x_{t})\mathop{}\!\mathrm{d}{t}+\sigma_{t}\,(x_{t}-\mu)\mathop{}\!\mathrm{d}{w}_{t}.$ (18) $$

To solve this SDE, we introduce a new variable $y_{t}=x_{t}-\mu$ and replace it into Eq. ([18](#A1.E18)), which typically yields a Geometric Brownian motion Ross [2014] on $y_{t}$, given by

$$ $\mathop{}\!\mathrm{d}{y}_{t}=-\theta_{t}y_{t}\mathop{}\!\mathrm{d}{t}+\sigma_{t}y_{t}\mathop{}\!\mathrm{d}{w}_{t}.$ (19) $$

To simplify the notation, we use $y$ rather than $y_{t}$ in all the following equations. This equation can be solved by applying Itô’s formula:

$$ $\displaystyle\begin{split}\mathop{}\!\mathrm{d}{\psi(y,t)}&=\frac{\partial\psi}{\partial t}(y,t)\mathop{}\!\mathrm{d}{t}+\frac{\partial\psi}{\partial y}(y,t)f(y,t)\mathop{}\!\mathrm{d}{t}\\ &\quad+\frac{1}{2}\frac{\partial^{2}\psi}{\partial y^{2}}(y,t)g(t)^{2}\mathop{}\!\mathrm{d}{t}\\ &\quad+\frac{\partial\psi}{\partial y}(y,t)g(t)\,\mathop{}\!\mathrm{d}{w},\end{split}$ (20) $$

where $\psi(y,t)=\ln|y|$ is a surrogate differentiable function. By substituting $f(y,t)$ and $g(t)$ with the drift and the diffusion functions in ([19](#A1.E19)), we obtain

$$ $\displaystyle\mathop{}\!\mathrm{d}{\psi(y,t)}=-(\theta_{t}+\frac{\sigma_{t}^{2}}{2})\mathop{}\!\mathrm{d}{t}+\sigma_{t}\mathop{}\!\mathrm{d}{w}.$ (21) $$

Then we can solve $y_{t}$ conditioned on $y_{s}$, by integrating both sides:

$$ $\ln|y_{t}|-\ln|y_{s}|=-\int_{s}^{t}(\theta_{z}+\frac{\sigma_{z}^{2}}{2})\mathop{}\!\mathrm{d}{z}+\int_{s}^{t}{\sigma}_{z}{\mathop{}\!\mathrm{d}{w}(z)}$ (22) $$

where the stochastic interaction follow a Gaussian distribution, i.e., $\int_{s}^{t}{\sigma}_{z}\mathop{}\!\mathrm{d}{w}(z)\sim\mathcal{N}\bigl{(}0,\int_{s}^{t}{\sigma}_{z}^{2}{\mathop{}\!\mathrm{d}{z}}\bigr{)}$, then we can rewrite:

$$ $\ln|y_{t}|=\ln|y_{s}|-(\bar{\theta}_{s:t}+\frac{\bar{\sigma}_{s:t}^{2}}{2})+\bar{\sigma}_{s:t}\epsilon_{s\rightarrow t},\quad\epsilon_{s\rightarrow t}\sim\mathcal{N}(0,I),$ (23) $$

where $\bar{\theta}_{s:t}=\int_{s}^{t}\theta_{z}\mathop{}\!\mathrm{d}{z}$, $\bar{\sigma}_{s:t}^{2}=\int_{s}^{t}\sigma_{z}^{2}\mathop{}\!\mathrm{d}{z}$, and $\bar{\sigma}_{s:t}=\sqrt{\bar{\sigma}_{s:t}^{2}}$. By replacing $y_{t}$ with the original $x_{t}-\mu$, we have the following

$$ $\ln|x_{t}-\mu|=\ln|x_{s}-\mu|-(\bar{\theta}_{s:t}+\frac{\bar{\sigma}_{s:t}^{2}}{2})+\bar{\sigma}_{s:t}\epsilon_{s\rightarrow t}.$ (24) $$

Note that the sign of $x_{t}-\mu$ remains consistent across all times $t$, therefore, we can safely omit the absolute value inside the logarithm, which leads to a log-normal distribution:

$$ $\log(\mu-x_{t})\sim\mathcal{N}\Bigl{(}\log(\mu-x_{s})-\int_{s}^{t}\bigl{(}\theta_{z}+\frac{1}{2}\sigma_{z}^{2}\bigr{)}\mathop{}\!\mathrm{d}z,\,\int_{s}^{t}\sigma_{z}^{2}\,\mathop{}\!\mathrm{d}z\,I\Bigr{)}.,$ (25) $$

which gives the Corollary [3.2](#S3.Thmtheorem2).
In addition, applying the exponential function to both sides yields

$$ $(x_{t}-\mu)=(x_{s}-\mu)\mathrm{e}^{-(\bar{\theta}_{s:t}+\frac{\bar{\sigma}_{s:t}^{2}}{2})+\bar{\sigma}_{s:t}\epsilon_{s\rightarrow t}}$ (26) $$

which is the solution to the SDE, and thus we complete the proof.
∎

###### Proposition A.1 .

###### Corollary A.2 .

###### Proof.

### A.2 Proof of the KL divergence

The KL divergence of ([9](#S3.E9)) can be derived as follows:

$$ $\displaystyle\tilde{L}$ $\displaystyle\coloneqq-\log p_{\phi}(x_{T})$ (27) $\displaystyle=-\log\int p_{\phi}(x_{0:T})\mathop{}\!\mathrm{d}x_{0:T-1}$ $\displaystyle=-\log\int\frac{p_{\phi}(x_{0:T})q(x_{0:T-1}\mid x_{T})}{q(x_{0:T-1}\mid x_{T})}\mathop{}\!\mathrm{d}x_{0:T-1}$ $\displaystyle=-\log\mathbb{E}_{q(x_{0:T-1}\mid x_{T})}\Bigl{[}\frac{p_{\phi}(x_{0:T})}{q(x_{0:T-1}\mid x_{T})}\Bigr{]}$ $\displaystyle\leq\underbrace{\mathbb{E}_{q(x_{0:T-1}\mid x_{T})}\Bigl{[}-\log\frac{p_{\phi}(x_{0:T})}{q(x_{0:T-1}\mid x_{T})}\Bigr{]}}_{\textit{negative evidence lower bound (ELBO)}}\quad\quad\text{(Jensen's Inequality)}$ $\displaystyle=\mathbb{E}_{q(x_{0:T-1}\mid x_{T})}\Bigl{[}-\log\frac{p(x_{0})\prod_{t=1}^{T}p_{\phi}(x_{t}\mid x_{t-1})}{\prod_{t=1}^{T}q(x_{t-1}\mid x_{t})}\Bigr{]}$ $\displaystyle=\mathbb{E}_{q(x_{0:T-1}\mid x_{T})}\Bigl{[}-\log p(x_{0})-\sum_{t=1}^{T}\log\frac{p_{\phi}(x_{t}\mid x_{t-1})}{q(x_{t-1}\mid x_{t})}\Bigr{]}$ $\displaystyle=\mathbb{E}_{q(x_{0:T-1}\mid x_{T})}\Bigl{[}-\log p(x_{0})-\sum_{t=1}^{T-1}\log\underbrace{\frac{p_{\phi}(x_{t}\mid x_{t-1})}{q(x_{t}\mid x_{t-1},x_{T})}\cdot\frac{q(x_{t}\mid x_{T})}{q(x_{t-1}\mid x_{T})}}_{\text{Bayes' rule on}\;q(x_{t-1}\mid x_{t},\,x_{T})}-\log\frac{p_{\phi}(x_{T}\mid x_{T-1})}{q(x_{T-1}\mid x_{T})}\Bigr{]}$ $\displaystyle=\mathbb{E}_{q(x_{0:T-1}\mid x_{T})}\Bigl{[}-\log p(x_{0})-\sum_{t=1}^{T-1}\log\frac{p_{\phi}(x_{t}\mid x_{t-1})}{q(x_{t}\mid x_{t-1},x_{T})}-\log\frac{q(x_{T-1}\mid x_{T})}{q(x_{0}\mid x_{T})}-\log\frac{p_{\phi}(x_{T}\mid x_{T-1})}{q(x_{T-1}\mid x_{T})}\Bigr{]}$ $\displaystyle=\mathbb{E}_{q(x_{0:T-1}\mid x_{T})}\Bigl{[}-\log\frac{p(x_{0})}{q(x_{0}\mid x_{T})}-\sum_{t=1}^{T-1}\log\frac{p_{\phi}(x_{t}\mid x_{t-1})}{q(x_{t}\mid x_{t-1},x_{T})}-\log p_{\phi}(x_{T}\mid x_{T-1})\Bigr{]}$ $\displaystyle=D_{KL}(q(x_{0}\mid x_{T})\mid\mid p(x_{0}))\quad\quad{(D_{KL}(p\|q)=\mathbb{E}[-\log\frac{p}{q}])}$ $\displaystyle\quad\quad+\sum_{t=1}^{T-1}D_{KL}(q(x_{t}\mid x_{t-1},x_{T})\mid\mid p_{\phi}(x_{t}\mid x_{t-1}))-\mathbb{E}_{q(x_{T-1}\mid x_{T})}\Bigl{[}\log p_{\phi}(x_{T}\mid x_{T-1})\bigr{]},$ $$

where the first term can be ignored since it doesn’t have trainable parameters, and the third term can be merged to the final stochastic flow matching objective.

### A.3 Derivation of Stochastic Flow Matching

#### Remark.

For two log-normal distributions with $\log p_{1}\sim\mathcal{N}(\mu_{1},\sigma_{1}^{2})$ and $\log p_{2}\sim\mathcal{N}(\mu_{2},\sigma_{2}^{2})$, the KL divergence between them is given by

$$ $D_{\text{KL}}(p_{1}\|p_{2})=\frac{(\mu_{1}-\mu_{2})^{2}}{2\sigma_{2}^{2}}+\frac{\sigma_{1}^{2}}{2\sigma_{2}^{2}}+\ln\frac{\sigma_{2}}{\sigma_{1}}-\frac{1}{2}.$ (28) $$

This result is for scalars but can be naturally extended to high-dimensional cases. In the following, we will use it to derive our stochastic flow matching objective.

More specifically, given the KL divergence

$$ $\mathbb{E}_{p}\Bigl{[}\,\sum_{t=0}^{T-1}D_{KL}(p(f_{\mu}(x_{t+1})\mid f_{\mu}(x_{t})\,\|\,p(f_{\phi}(x_{t+1},t)\mid f_{\phi}(x_{t},t))\,\Bigr{]}$ (29) $$

and the truth that $f_{\phi}(x_{t},t)$ approximates $\mu-x_{t}$. Since the two transitions are log-normally distributed and share the same parameters $\theta_{t}$ and $\sigma_{t}$ as in Eq. ([7](#S3.E7)), we can initially obtain a log space loss based on the Remark ([28](#A1.E28)), as

$$ $L\coloneqq\mathbb{E}_{\mu\sim p_{\text{data}},x_{t}\sim q(x_{t}|x_{0},\mu)}\Bigl{[}\,\frac{1}{2\sigma_{t+1}^{2}}\|\log f_{\mu}(x_{t},t)-\log f_{\phi}(x_{t},t)\|^{2}\,\Bigr{]}.$ (30) $$

However, this would run into issues when $\mu-x_{t}\leq 0$. Though this can be alleviated using absolute values and adding a small additive term $\epsilon$, it complicates the learning and is still unstable.

To address it, we assume that, after some training, $f_{\phi}(x_{t},t)/f_{\mu}(x_{t},t)=1+\delta$ where $|\delta|\ll 1$. The first-order Taylor approximation is then

$$ $\log f_{\phi}(x_{t},t)-\log f_{\mu}(x_{t},t)=\log(1+\delta)\approx\frac{f_{\phi}(x_{t},t)-f_{\mu}(x_{t},t)}{f_{\mu}(x_{t},t)}.$ (31) $$

Since the denominator does not depend on the parameters,
we obtain our simplified loss function by omitting all non-trainable weights:

$$ $L\coloneqq\mathbb{E}_{\mu\sim p_{\text{data}},x_{t}\sim q(x_{t}|x_{0},\mu)}\Bigl{[}\,\|f_{\mu}(x_{t},t)-f_{\phi}(x_{t},t)\|^{2}\,\Bigr{]},$ (32) $$

which is the proposed stochastic flow matching objective.

## Appendix B Maximum Likelihood Estimation

Following DDPMs Ho et al. [2020], both our training and sampling are implemented with discrete times, which can of course be converted into continuous times, as in Score SDEs Song et al. [2021], but that requires $\theta$ and $\sigma$ schedules to be integrable. Below, we further show that the solution to FoD also allows us to compute the maximum likelihood:

Given the clean data $\mu$, assuming that there exists an optimal forward transition from $z_{t}$ to $z_{t+1}$, where $z_{t}=|\mu-x_{t}|$. In other words, we want to maximize the likelihood of $p(z_{t+1}\mid z_{t})$, which is a log-normal distribution as illustrated in Corollary 3.2, and its density is given by

$$ $p(z_{t+1}\mid z_{t})=\frac{1}{z_{t+1}\sigma_{t+1}\sqrt{2\pi}}\exp\left({-\frac{1}{2\sigma_{t+1}^{2}}\left[\ln z_{t+1}-lnz_{t}+\big{(}\theta_{t+1}+\frac{\sigma_{t+1}^{2}}{2}\big{)}\right]^{2}}\right).$ (33) $$

Then, we can minimize the negative log-likelihood:

$$ $\displaystyle\begin{split}-\ln p(z_{t+1}\mid z_{t})&=\ln z_{t+1}+\frac{1}{2\sigma_{t+1}^{2}}\left[\ln z_{t+1}-lnz_{t}+\big{(}\theta_{t+1}+\frac{\sigma_{t+1}^{2}}{2}\big{)}\right]^{2}+\cancel{\ln(\sigma_{t+1}\sqrt{2\pi})},\end{split}$ (34) $$

which can be solved by setting the gradient to 0, as

$$ $\displaystyle\begin{split}\nabla_{z_{t+1}}-\ln p(z_{t+1}\mid z_{t})&=\frac{1}{z_{t+1}}\left[1+\frac{1}{\sigma_{t+1}^{2}}\left(\ln z_{t+1}-lnz_{t}+\big{(}\theta_{t+1}+\frac{\sigma_{t+1}^{2}}{2}\big{)}\right)\right]=0.\end{split}$ (35) $$

Since $z_{t+1}$ is not zero, the optimal solution $z_{t+1}^{*}$ is obtained according to

$$ $z_{t+1}=z_{t}\mathrm{e}^{-(\theta_{t+1}+\frac{\sigma_{t+1}^{2}}{2})-\sigma_{t+1}^{2}}.$ (36) $$

Recall that $\mu-x_{t}$ has the same sign for all times $t$. Replacing $z_{t}$ with $|\mu-x_{t}|$ gives the following:

$$ $(\mu-x_{t+1})^{*}=(\mu-x_{t})\mathrm{e}^{-(\theta_{t+1}+\frac{\sigma_{t+1}^{2}}{2})-\sigma_{t+1}^{2}},$ (37) $$

which is the optimal forward flow from $x_{t+1}$ to $\mu$.

Based on it, we can get the maximum likelihood learning objective:

$$ $L=\big{\|}(\mu-x_{t+1})^{*}-\mathbb{E}[\mu_{\phi}-x_{t+1}]\big{\|}^{2},$ (38) $$

$$ $L=\big{\|}x_{t+1}^{*}-\mathbb{E}_{\phi}[x_{t+1}\mid x_{t}]\big{\|}^{2},$ (39) $$

where $\mathbb{E}_{\phi}[x_{t+1}\mid x_{t}]$ is the expectation given $x_{t}$ in discrete time: $\mathbb{E}_{\phi}[x_{t+1}\mid x_{t}]=x_{t}+\mathop{}\!\mathrm{d}{x}_{t}$.

$$ $\displaystyle\begin{split}\mathbb{E}\big{[}\mu_{\phi}-x_{t+1}\big{]}&=(\mu_{\phi}-x_{t})\mathrm{e}^{-(\theta_{t+1}+\frac{\sigma_{t+1}^{2}}{2})+\frac{\sigma_{t+1}^{2}}{2}}\\ &=(\mu_{\phi}-x_{t})\mathrm{e}^{-\theta_{t+1}}.\end{split}$ (40) $$

Combining ([37](#A2.E37)) and ([40](#A2.E40)) predicts $x_{t+1}$ in the optimal path. This maximum likelihood-based loss function performs similarly to stochastic flow matching and can be potentially used in future work.

## Appendix C More Experimental Details

### C.1 Implementation

We set the number of diffusion steps to 100 for all tasks. Practically, we choose to use a discrete time implementation for our method, where we let $\bar{\theta}_{t}=\int_{0}^{t}\theta_{z}\mathop{}\!\mathrm{d}z\approx\sum\theta_{t}\Delta t$ and $\bar{\sigma}_{t}^{2}=\int_{s}^{t}\sigma_{z}^{2}\,\mathop{}\!\mathrm{d}z\approx\sum\sigma_{t}^{2}\Delta t$. To ensure the FoD converges to the clean data $\mu$, we let the deterministic exponential term at the terminal state be a smaller value i.e. $\mathrm{e}^{-\int_{0}^{t}(\theta_{s}+\frac{1}{2}\sigma_{s}^{2})\mathop{}\!\mathrm{d}{s}}=\delta=0.001$. Solving it leads to an updated time interval $\Delta t=\frac{\log\delta}{\int_{0}^{t}(\theta_{s}+\frac{1}{2}\sigma_{s}^{2})\mathop{}\!\mathrm{d}{s}}$.

### C.2 Image Restoration Datasets

The details of the image restoration tasks and their datasets are listed below:

- •
Deraining: collected from the Rain100H [Yang et al., 2017] dataset containing 1800 images for training and 100 images for testing.
- •
Dehazing: collected from the RESIDE-6k [Qin et al., 2020] dataset which has mixed indoor and outdoor images with 6000 images for training and 1000 images for testing.
- •
Low-light enhancement: collected from the LOL [Wei et al., 2018] dataset containing 485 images for training and 15 images for testing.
- •
Face inpainting: we use CelebaHQ as the training dataset and divide 100 images with 100 thin masks from RePaint [Lugmayr et al., 2022] for testing.

### C.3 Additional Discussions

Below, we discuss faster sampling of our model and the FoD process for each task.

As mentioned in Section 3.3, we provide two fast sampling strategies with Markov and non-Markov Chains. Their comparison on image deraining is illustrated in Section 5.1 of the main paper. Here, we give more comparison results on four image restoration tasks, reporting PSNR, SSIM, LPIPS, and FID values. The results with standard FoD sampling (See Algorithm [2](#alg2)) are also reported as the baseline. It can be observed that, in most tasks, decreasing the number of steps in fast sampling approaches leads to better performance, particularly in terms of PSNR and SSIM. Their perceptual performance also decreases when the number of steps decreases from 20 to 5, which suggests we set the sample step to 10 as a practical rule of thumb for efficient sampling. Note that the key difference between faster sampling approaches and the standard FoD sampling is that the latter’s iterative process highly depends on the $\mu$ prediction, while the former gradually refines the state $x_{t}$ rather than $\mu$. Also, we would recommend using non-Markov chain sampling for small degradation problems (which means $\mu$ prediction is easier), such as image deraining and dehazing. The visual comparisons are illustrated in Figure [8](#A3.F8).

We further provide more examples of the visualization of the FoD sampling process in Figure [9](#A3.F9). It can be observed that, for image restoration tasks, the noise level for each task and even for each image is different. This is given by the term $\mu-x_{0}$ in the solution, where areas with large difference (between LQ and HQ images) tend to produce large noise.

### C.4 Additional Results

In this section, we provide more results for four image restoration tasks in including image deraining, low-light enhancement, image dehazing, and image inpainting in Figure [10](#A3.F10), Figure [11](#A3.F11), Figure [12](#A3.F12), and Figure [13](#A3.F13). In most tasks, the results produced by our method are sharper and more realistic. In addition, we also provide more results for unconditional generation with SDE and ODE on CIFAR-10 in Figure [14](#A3.F14) and Figure [15](#A3.F15), respectively. Please zoom in
for the best view.

Figure: Figure 7: Comparison of different sampling approaches with pretrained FoD models on four image restoration tasks. The baseline is the Euler-Maruyama method with 100 sampling steps.
Refer to caption: /html/2505.16733/assets/x14.png

Figure: Figure 8: Comparison of fast sampling with Markov and non-Markov chains with different steps.
Refer to caption: /html/2505.16733/assets/x18.png

Figure: Figure 9: Visualization of the diffusion process using trained FoD models on various tasks, including deraining, dehazing, low-light enhancement, face inpainting, and unconditional generation. In each case, FoD gradually injects noise into the degraded regions and subsequently denoises these intermediate states, restoring images with enhanced and corrected details.
Refer to caption: /html/2505.16733/assets/x19.png

Figure: Figure 10: Visual results of image deraining on Rain100H Yang et al. [2017] dataset.
Refer to caption: /html/2505.16733/assets/x20.png

Figure: Figure 11: Visual results of image low-light enhancement on LOL Wei et al. [2018] dataset.
Refer to caption: /html/2505.16733/assets/x21.png

Figure: Figure 12: Visual results of image dehazing on RESIDE-6k Qin et al. [2020] dataset.
Refer to caption: /html/2505.16733/assets/x22.png

Figure: Figure 13: Visual results of image inpainting on CelebA-HQ Karras et al. [2017] dataset.
Refer to caption: /html/2505.16733/assets/x23.png

Figure: Figure 14: Unconditional generation by FoD (SDE sampler).
Refer to caption: /html/2505.16733/assets/supp/supp-sde_sample.png

Figure: Figure 15: Unconditional generation by the ODE variant of FoD.
Refer to caption: /html/2505.16733/assets/supp/supp-ode_sample.png

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 4
questions: **(Q1)**: Given the primary weakness is unconditional generation, have the authors experimented with modifying the prior distribution to better match the model's log-normal structure (e.g., starting from a log-normal prior)

**(Q2)**: Do the authors have explanations for the behavior between the MC and non-MC samplers in Figure 4? Why are structural metrics better here compared to generation quality metrics? How should one decide between the two samplers?

**(Q3)**: What is $x_s$ and $\mu$ in the tasks outlined in Figure 3? How does the SDE behave when the tasks represents bridging two very semantically different distributions?

**(Q4)**: How sensitive is the model to the choice of the $\sigma_t$ schedule given that it now controls mean-reversion / a state dependent term in the diffusion term?
soundness: 4
strengths: **(S1)**: Elegant formulation. A single, forward-only SDE for diffusion with the mean-reversion term is a neat formulation for image-restoration and conditional generation tasks. The mean-reversion term enables a state-dependent denoising process that dynamically adjusts to different corruption levels within an image. To me, this makes a lot of sense for conditional generation.

**(S2)**: Tractability and flexibility. The SDE with mean-reversion in the diffusion term still yields a unique, tractable solution and enables a simple loss function that is stable to train. This is a strong point in favor of the method. Providing both Markovian and non-Markovian sampling strategies for this forward-only SDE further improves the flexibility of this approach.

**(S3)**: Good ablations. Nullifying the diffusion term and reverting the SDE back to the flow-matching ODE clearly demonstrates worse performance on image restoration tasks, in terms of structural similarity metrics. This demonstrates the need for stochasticity in the mean-reverting SDE. Another ablation on fast-sampling is helpful. 

Overall, I think the idea and application of this paper is novel and interesting, and so I would recommend it for acceptance. If some of my concerns outlined below are addressed, I would be happy to raise my score.
summary: This paper introduces forward-only diffusion, which replaces the standard diffusion SDE with a mean-reversion term in both the drift and diffusion terms. They derive a tractable solution to the SDE that determines the conditional distribution $p(x_{t+1}|x_t, \mu)$ and parameterize their neural network to learn the flow $\hat{\mu}_{\phi} - x_t$. This enables a tractable loss to minimize the KL between the ground truth conditional and the model's estimate. Experimental results on image restoration tasks are provided, demonstrating improvements over baselines.
weaknesses: **(W1)**: Poor unconditional generation. The FID scores on CIFAR-10 (7.89 for FoD-SDE, 5.01 for FoD-ODE) are not competitive with standard forward-backward diffusion models (e.g., Score SDE @ 2.38) or even other forward-only ODE models like Rectified Flow (2.58). While noted as a limitation, this positions FoD as more a specialized method for conditional image generation tasks than as a general generative model. 

**(W2)**: Limited exploration of conditional image generation tasks. The paper mainly focuses on image restoration, which is a low-entropy task (i.e. the source is already close to the target image). Some qualitative examples for image-to-image translation are provided, but a more extensive evaluation on translation tasks would be useful to support the generality of this approach. Additional experiments on text-to-image or latent-diffusion architectures would be further welcome for completeness.

**(W3)**: Missing comparisons. Some recent work on diffusion bridges (denoising diffusion bridge models) tackle similar problems as this paper. Comparison and a more detailed discussion around bridge models is missing. While FoD is an instantiation of stochastic interpolant methods, the paper could be strengthened by a more detailed comparison to other recent SI-based methods that have also been applied to image restoration.

---

confidence: 3
contribution: 2
initial_rating: 4
presentation: 3
questions: N/A
soundness: 2
strengths: * The writing is easy to follow.
* Using the proposed mean-reverting-style SDE for generative modeling looks novel to me.
summary: This manuscript proposes a new forward diffusion process for generative modeling. Specifically, the authors leverage a mean-reverting style SDE and apply the mean-reversion terms to both the drift and diffusion coefficients. The designed SDE has a closed-form solution, and the only unknown term is the simulation target. Following the spirit of denoising score matching, the author proposes a simple regression objective to learn it and thus enable simulating the designed SDE. For the sampling option, the author explored both the Euler method and the first-order discretization. Experimental results demonstrate that the proposed technique is applicable on image-to-image tasks.
weaknesses: * The claim about "simpler, single" diffusion process is somewhat unconvincing to me. According to the training algorithm and sampling procedure, the effort is almost the same as that of the diffusion models, and the training objective itself needs approximation.
* Another main claim of the paper is that the proposed method can be viewed as a stochastic counterpart to flow matching. However, to me, it would be necessary to compare the established stochastic counterpart, known as diffusion bridges or bridge matching [1, 2], of the flow matching model, both conceptually and empirically (on image-to-image benchmarks).
* For empirical results, it would be better to add Gassuain-to-image generation tasks to demonstrate the effectiveness of the proposed framework.



## References
[1] Peluchetti, Stefano. ‘Non-Denoising Forward-Time Diffusions’. (2023)

[2] Shi, Yuyang, et al. ‘Diffusion Schrödinger Bridge Matching’. (NeurIPS 2023)

---

confidence: 4
contribution: 1
initial_rating: 2
presentation: 3
questions: 1. The authors need to discuss more how selection of the coefficients impact the behavior of the diffusion. 
2. Choice of $\theta$ and $\sigma$ affects the performance? 
3. In P4, line 146, the authors stated that As $t \rightarrow \infty$, the SDE converges to a stationary state $x_T ∼ N(x_t | \mu, \lambda^2)$. This formula does not make sense to me. Please check if t is T.
soundness: 2
strengths: The paper is well written. The authors effort to consider diffusion coefficients as state dependent although the implementation is weak.
summary: This paper proposes a probabilistic forward-only diffusion model (FoD) that incorporates Geometric Brownian Motion (GBM) to transform an image into another that is in two distributions separately. The authors claimed that by introducing a mean-reversion term $\mu - x_t$​ into both the drift and diffusion coefficients of the SDE, they defined an analytically solvable forward-only process, eliminating the need to approximate or learn a reverse SDE. However, I have several concerned. First, I believe this claim of FoD is actually a backward diffusion (as I explained in Weakness section).  Second, in conventional diffusion models, the process involves diffusing an image  to noise (forward), and then reconstructing from noise to an image that closely remember the original image (reverse) although one can always  transform between two deterministic images. The setup in this paper actually resembles the GOUB., more recently; UNIDB. Although the authors proposed to  add the state dependency into the diffusion coefficient, there is no proof of the effect of this diffusion. Indeed the author almost has not discussion on the coefficients. Instead, the authors used extremely small diffusion coefficients, this made it actually more close to deterministic method, reducing the diffusion effects, which I believe making this paper relatively weak contribution to the field.
weaknesses: 1. It seems the authors misinterpretation of forward diffusion as backward. This can be observed from they actually started from $x_T$ ∼ $p_{data}$ to  $x_0$ ∼$ p_{prior}$. The training step stated in the paper actually resembles the backward diffusion diffusion step in the conventional diffusion. There is not much different between this method in terms forward/backward as the diffusion itself is not a reason to stay on forward only as what we can see from GOUB.

2. P4. Line 178, as the author mentioned, "The subtractive form of the logarithm reflects that
the flow field decays multiplicatively from its initial value with a stochastic exponential scaling", this translate to the original image difference, the variance of the noise can explore, I wonder how the authors address this issue. The only solution in the paper is that they set the $ e^{-\int_0^t (\theta_s+\frac{1}{2} \sigma_s^2) ds}= 0.001$. The authors should have discussed the impact of selecting such small diffusion coefficients. This small value seems to make it almost a deterministic process. 

3. As I mentioned in the summary and item 2, the authors used extremely small diffusion coefficients, this made it actually more close to deterministic method, reducing the diffusion effects. This can be observed in several sequences of images. For example, the second row images (over various time steps) in Figure 1 seemed so deterministic.

4. The following paper should be referred:
1). K Zhu et. al. UniDB: A Unified Diffusion Bridge Framework via Stochastic Optimal Control
2). G Kim et. al. Diffusion-based generative model for financial time series via geometric Brownian motion

5.  PROPOSITION 3.1 seems like an existing result and no need to prove again.
6. Although the paper compared their method to GOUB in terms of metrics such as FID, they did not provide the GOUB results in the figures of images that compares the results. It will be good to seem their comparisons in images as they are similar diffusions between two deterministic images.


Decision: Reject