MaRS: A Fast Sampler for Mean Reverting Diffusion based on ODE and SDE Solvers

We really appreciate the thoughtful and constructive comments of the reviewer. Below we respond to the weaknesses and questions.

Q1(a)&W1: How does reducing NFEs affect sampling quality?

A: From the perspective of reverse-time SDE and PF-ODE, the estimation error of the solution plays a critical role in determining the sampling quality. This estimation error is primarily influenced by $h_{max} = \max_{1 \leq i \leq M}{\lambda_i - \lambda_{i-1}} = \mathcal{O}(\frac{1}{M})$, where $M$ represents the NFE. During sampling, the noise schedule must be designed in advance, which also determines the $\lambda$ schedule. Since $\lambda_i$ is monotonically increasing, a larger NFE results in a smaller $h_{max}$, thereby reducing the estimation error and improving sampling quality. However, different sampling algorithms exhibit varying convergence rates. Experimental results show that the MR Sampler (our method) can achieve a good score with as few as 5 NFEs, whereas posterior sampling and Euler discretization fail to converge even with 50 NFEs.

Specifically, we conducted additional experiments on the desnow task with low NFE settings, and the results are presented in Table below. The sampling quality remains largely stable for NFE values larger than 10, gradually deteriorates when the NFE is below 10, and collapses entirely when NFE is reduced to 2.

MR Sampler-SDE-2 with data prediction and uniform $\lambda$ on the desnow task

Q2: Do different image restoration tasks benefit differently from the MR Sampler in terms of speed or quality?

A: We consider that different image restoration tasks do not benefit differently from the MR Sampler in terms of speed or quality. The performance of the MR Sampler varies across different tasks; however, this variation originates from the neural network rather than the sampling algorithm itself. Our method is a sampling acceleration algorithm specifically designed for MR Diffusion and does not incorporate any task-specific prior knowledge in its derivation. From this perspective, the nature of the task does not inherently affect the performance of MR Sampler. However, our derivation involves performing the Taylor expansion on the neural network function. As a result, the performance of our algorithm is inevitably influenced by the characteristics and effectiveness of the neural network.

Q3: Why is Mean-Reverting SDE better for posterior sampling compared to other methods like Optimal Transport or Schrödinger Bridge?

A: Both Schrödinger Bridge and Optimal Transport approaches arrive at conclusions similar to those of MR Diffusion, albeit through different methodologies:

$$q(X_t|X_0,X_1)=\mathcal{N}(X_t;\mu_t(X_0,X_1),\Sigma_t), \\quad \mu_t=\frac{\bar\sigma_t^2}{\bar\sigma_t^2+\sigma_t^2}X_0+\frac{\sigma_t^2}{\bar\sigma_t^2+\sigma_t^2}X_1, \\;\Sigma_t=\frac{\bar\sigma_t^2\sigma_t^2}{\bar\sigma_t^2+\sigma_t^2}I \tag{1}$$ $$x_t=\mu+(x_s-\mu)e^{-\bar\theta_{s:t}}+\lambda^2(1-e^{-2\bar\theta_{s:t}})\epsilon \tag{2}$$

Equation (1) is presented in Proposition 3.3 of paper [1], while Equation (2) is given in Proposition 3.1 of paper [2]. Despite differences in notation and definitions, the principles described by the two equations are fundamentally consistent. However, a key distinction lies in how the posterior sampling is implemented. Schrödinger Bridge adopts the posterior sampling method of DDPM without leveraging the information of $X_N$ (see Algorithm 2 in paper [1] for further details). In contrast, MR Diffusion incorporates the information of $\mu$ in its posterior sampling (refer to Equation (2) for details), which leads to improved performance.

[1]. Liu, G. H., Vahdat, A., Huang, D. A., Theodorou, E. A., Nie, W., & Anandkumar, A. (2023). I $^ 2$ SB: Image-to-Image Schrödinger Bridge. arXiv preprint arXiv:2302.05872.

[2]. Ziwei Luo, Fredrik K Gustafsson, Zheng Zhao, Jens Sjölund, and Thomas B Schön. Image restoration with mean-reverting stochastic differential equations. arXiv preprint arXiv:2301.11699, 2023b.

Thank you for the interest and acknowledgment of our contributions and the insightful questions. Below we respond to the weaknesses and questions.

Q1&W1: How to determine the optimal number of sampling steps

A: The NFE is strictly positively correlated with sampling quality. From the perspective of reverse-time SDE and PF-ODE, the estimation error of the solution plays a critical role in determining the sampling quality. This estimation error is primarily influenced by $h_{max} = \max_{1 \leq i \leq M}{\lambda_i - \lambda_{i-1}} = \mathcal{O}(\frac{1}{M})$, where $M$ represents the NFE. During sampling, the noise schedule must be designed in advance, which also determines the $\lambda$ schedule. Since $\lambda_i$ is monotonically increasing, a larger NFE results in a smaller $h_{max}$, thereby reducing the estimation error and improving sampling quality.

However, different sampling algorithms exhibit varying convergence rates. Experimental results show that the MR Sampler (our method) can achieve a good score with as few as 5 NFEs, whereas posterior sampling and Euler discretization fail to converge even with 50 NFEs.

Specifically, we conducted additional experiments on the desnow task with low NFE settings, and the results are presented in Table below. The sampling quality remains largely stable for NFE values larger than 10, gradually deteriorates when the NFE is below 10, and collapses entirely when NFE is reduced to 2. Based on our experience, we recommend using 10–20 NFEs, which provides a reasonable trade-off between efficiency and performance. We have added a section to appendix to provide an in-depth discussion on determining the NFE.

MR Sampler-SDE-2 with data prediction and uniform $\lambda$ on the desnow task

Q2: What was the rationale behind choosing a low-light and motion-blurry dataset for the visualizations in Figure 2?

A: We conducted extensive experiments on ten datasets. Due to the page limitations of the main text, it was not feasible to include results for all datasets. As a result, we randomly selected two datasets to serve as representatives, and the results for the remaining datasets are provided in Appendix D2.

Q3: In reference to the related paper [1], it is mentioned that solving the SDE typically requires only 22 steps, although they still use 100 steps. Since you opted to use 100 steps aswell, I am wondering why the performance with 20 steps is still so substantially different.

A: We would like to emphasize that in paper [1], the authors were able to achieve sampling with 22 NFEs only for denoising tasks. In the sampling process of MR Diffusion, the initial state is $\mu + \sigma\epsilon$, where $\mu$ is typically set to a low-quality image, and the final state is $\boldsymbol{x}_0$, which is usually a high-quality image. For denoising tasks, the authors assume that the low-quality image is equivalent to the high-quality image plus Gaussian noise, implying $\boldsymbol{x}_0 = \mu$ (see Section 4.3 and Appendix B in paper [1]). This leads to $\boldsymbol{x}\_t = \mu + \sigma\epsilon = \boldsymbol{x}\_0 + \sigma\epsilon$. In this case, the PF-ODE simplifies to $\mathrm{d}\boldsymbol{x}=-\frac12\boldsymbol{\sigma}^2\_t\mathrm{e}^{-2\bar\theta\_t}\nabla\_{x}\log{p\_t(\boldsymbol{x})}\mathrm{d}t.(1)$

It is important to note that this scenario is equivalent to setting $f(x,t) = 0$, causing the linear term in the PF-ODE to cancel out. In our paper, we explain in Section 3.3 the reasons for the slow convergence of Euler discretization, which introduces approximation errors from both linear and nonlinear terms. However, for Equation (1), the absence of the linear term eliminates the corresponding error, allowing for the use of a smaller NFE in this specific case. Such simplification may not hold for other tasks because they require the degradation of the image to result exclusively from additive Gaussian noise. In contrast, our sampling algorithm is not subject to this limitation, making it more broadly applicable.

[1]. Ziwei Luo, Fredrik K Gustafsson, Zheng Zhao, Jens Sjölund, and Thomas B Schön. Image restoration with mean-reverting stochastic differential equations. arXiv preprint arXiv:2301.11699, 2023b.

We thank the reviewer for the thoughtful and constructive comments. Below we provide detailed responses to the questions.

Q1: Why wasn't an ablation study done on n, k?

A: Because the variable $n$ is not explicitly defined or utilized in our paper. We assume that the reviewer uses $n$ to refer to the NFE (number of function evaluations) and $k$ as the order of the algorithm. We have performed ablation studies on values of NFE and $k$. Detailed results are shown in Figure 2 of the main text and Tables 6–15 in Appendix D2.

Q2: Why wasn't wall clock time reported? Please provide.

A: We regret not reporting the wall clock time in our original manuscript. We indeed measured it in our experiments on a single NVIDIA A800 GPU. The average wall clock time per image of two representative tasks (low-light and motion-blurry image restoration) are reported below, and we have added a section D.5 to the Appendix to report the wall clock time.

As observed, the time consumption and NFE are approximately proportional. This relationship arises because the sampling time is predominantly determined by neural network inference, while the time required for other computations is negligible in comparison.

Q3: comparison to standard (non-MR) SDE fast samplers

A: To the best of our knowledge, existing sampling acceleration methods are specifically designed for DPM, so these methods are not directly applicable to MR Diffusion. Therefore, when comparing the MR Sampler and the standard SDE Sampler, the diffusion models used by the two methods are different. Specifically, the standard SDE Sampler is for the DPM, while the MR Sampler is for the MR Diffusion. However, to compare standard SDE sampler with MR Sampler is questionable, as DPM and MR Diffusion exhibit different levels of performance on the image restoration task. Under such conditions, it becomes difficult to discern whether the observed performance differences arise from the sampling algorithm itself or from the diffusion model.

Q4: How was guidance incorporated for image restoration? Does the particular method for incorporating the degraded observation make a difference for the performance of MR sampler?

A: MR Diffusion (MRD) is indeed a type of conditional diffusion model. Luo et al. [1] developed MRD by redesigning the drift coefficient of the SDE and incorporating $\mu$ into the diffusion model. In the diffusion process of DPM, the trajectory starts from $\boldsymbol{x}_0$ and converges to pure Gaussian noise. In contrast, the diffusion process of MRD begins at $\boldsymbol{x}_0$ and converges to $\mu + \sigma\epsilon$. The introduction of $\mu$ serves as a conditional input channel for MRD. For image-to-image tasks, MRD assigns the input image to $\mu$, eliminating the need to incorporate image conditions via Classifier-Free Guidance (CFG).

Actually, incorporating $\mu$ does not make a favorable difference to the performance of MR Sampler. However, it is worth noting that the inclusion of $\mu$ increases the complexity of the reverse-time SDE and PF-ODE, making the solution of these equations more challenging.

[1]. Ziwei Luo, Fredrik K Gustafsson, Zheng Zhao, Jens Sjölund, and Thomas B Schön. Image restoration with mean-reverting stochastic differential equations. arXiv preprint arXiv:2301.11699, 2023b.

We thank the reviewer for the thoughtful and constructive comments. Below we provide detailed responses to the questions.

W1: reason for using the backward difference method

A: We use the backward difference method to estimate the derivative of the neural network. Here we take Equation (14) from our paper as an example.

$$-2\alpha\_t\int\_{\lambda_s}^{\lambda_t}e^{-\lambda}\boldsymbol{\epsilon}\_\theta(x\_\lambda,\lambda)\mathrm{d}\lambda= -2\sigma\_{t}\sum\_{n=0}^{k-1}\left[\boldsymbol{\epsilon}\_\theta^{(n)}(\boldsymbol{x}\_{\lambda_{s}},\lambda_{s})\left(e^{h}-\sum\_{m=0}^n\frac{(h)^m}{m!}\right)\right] +\mathcal{O}(h^{k+1}) \tag{14}$$

Equation (14) represents the non-linear component of the solution presented in Equation (13). Our objective is to estimate the derivative of the neural network $\boldsymbol{\epsilon}^{(n)}\_\theta(\boldsymbol{x}\_{\lambda},\lambda)$ with respect to the signal-to-noise ratio (SNR) $\lambda$. Using $\boldsymbol{\epsilon}^{(1)}\_\theta(\boldsymbol{x}\_{\lambda},\lambda)$ as a specific example, as illustrated in Algorithms 1–6, this process involves discretizing time and computing the solution iteratively. Specifically, we need to estimate $\boldsymbol{\epsilon}^{(1)}\_\theta(\boldsymbol{x}\_{\lambda_i},\lambda_i)$ at each time step. Because the analytical expression of $\boldsymbol{\epsilon}\_\theta$ is unavailable, the derivative is commonly approximated by the difference method, including forward difference, central difference, and backward difference approaches. However, the forward and central difference methods require the value of $\boldsymbol{\epsilon}\_\theta(\boldsymbol{x}\_{\lambda_{i+1}},\lambda_{i+1})$, but $\boldsymbol{x}\_{\lambda_{i+1}}$ itself needs to be predicted. Consequently, we adopt the backward difference method in this case.

W2: current SOTA methods

A: To the best of our knowledge, existing sampling acceleration methods are specifically designed for DPM, and our method is the first sampling acceleration algorithm for MR Diffusion. The corresponding SDE for DPM is given as $\mathrm{d}\boldsymbol{x}=f(t)\boldsymbol{x}\mathrm{d}t + g(t) \mathrm{d}\boldsymbol{w} \;(1)$, whereas the SDE for MR Diffusion is $\mathrm{d}\boldsymbol{x}=f(t)\left(\boldsymbol{\mu}-\boldsymbol{x}\right)\mathrm{d}t+g(t)\mathrm{d}\boldsymbol{w} \;(2)$. When $\boldsymbol\mu=0$, it can be shown that Equation (2) reduces to Equation (1). This implies that DPM is a special case of MR Diffusion, and MR Diffusion can be regarded as a generalization of DPM. Consequently, existing fast sampling methods designed for DPM are not directly applicable to MR Diffusion.

W3: "frequently fall outside" in line304

A: The original text is "During the training phase, the noise prediction neural network is designed to fit normally distributed Gaussian noise. Although the standard deviation of this Gaussian noise is set to 1, the values of samples can frequently fall outside the range of [-1, 1]." We apologize for any confusion here. The original statement pertains to the $3\sigma$ rule of the Gaussian distribution. Specifically, for a random variable $x$ that follows a normal distribution, the probability of $x \in [-1, 1]$ is approximately 65.26%, while the probability of $x \notin [-1, 1]$ is 34.74%. We acknowledge that the use of frequently in this context is not rigorous. We will revise the text as follows: "When the standard deviation of this Gaussian noise is set to 1, the values of samples can fall outside the range of [-1,1] with a probability of 34.74%."

W4: the neural network architecture used in experiments

A: We sincerely apologize for not explicitly clarifying the model framework and neural network employed in our experiments. Our work fully adheres to the framework of DA-CLIP [1], an image restoration model designed to address multiple degradation problems simultaneously without requiring prior knowledge of the degradation. The diffusion model in DA-CLIP is derived from Refusion [2], and its neural network architecture is based on NAFNet. NAFNet builds upon the U-Net architecture by replacing traditional activation functions with SimpleGate and incorporating an additional multi-layer perceptron to manage channel scaling and offset parameters for embedding temporal information into the attention and feedforward layers. For further details, please refer to Section 4.2 in [2]. We have added a section D.2 to the Appendix to describe the network architecture used in experiments.

[1]. 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.

[2]. Ziwei Luo et al. 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, pp. 1680–1691, 2023c.