Come-Closer-Diffuse-Faster: Accelerating Conditional Diffusion Models for Inverse Problems through Stochastic Contraction

Introduction

Denoising diffusion models and score-based models are new trending classes of generative models, which have recently drawn signficant attention amongst the community due to their state-of-the-art performance. Although inspired differently, both classes share very similar aspects, and can be cast as variants of each other , thus they are often called diffusion models.

In the forward diffusion process, a sampled data point ${\bm{x}}$ at time $t=0$ is perturbed gradually with Gaussian noise until $t=T$, arriving approximately at spherical Gaussian distribution, which is easy to sample from. In the reverse diffusion process, starting from the sampled noise at $t=T$, one uses the trained score function to gradually denoise the data up to $t=0$, arriving at a high quality data sample.

Interestingly, diffusion models can go beyond unconditional image synthesis, and have been applied to conditional image generation, including super-resolution , inpainting , MRI reconstruction , image translation , and so on. One line of works re-design the diffusion model specifically suitable for the task at hand, thereby achieving remarkable performance on the given task . However, they compromise flexibility since the model cannot be used on other tasks. Another line of works, on which we build our method on, keep the training procedure intact, and only modify the inference procedure such that one can sample from a conditional distribution . These methods can be thought of as leveraging the learnt score function as a generative prior of the data distribution, and can be flexibly used across different tasks.

Unfortunately, a critical drawback of diffusion models is that they are very slow to sample from. To address this, for unconditional generative models, many works focused on either constructing deterministic sample paths from the stochastic counterparts , searching for the optimal steps to take after the training of the score function , or by retraining student networks that can take shortcuts via knowledge distillation . Orthogonal and complementary to these prior works, in this work, we focus on accelerating conditional diffusion models by studying the contraction property of the reverse diffusion path.

Specifically, our method, which we call Come-Closer-Diffuse-Faster (CCDF), first perturbs the initial estimate via forward diffusion path up to $t_{0}<T$, where $t_{0}$ denotes the time where the reverse diffusion starts. This forward diffusion comes almost for free, without requiring any passes through the neural network. While the distribution of forward-diffused (noise-added) images increases the estimation errors from the initialization as shown in Fig. 2(b), the key idea of the proposed CCDF is that the reverse conditional diffusion path reduces the error exponentially fast thanks to the contraction property of the stochastic difference equation . Therefore, compared to the standard approach that starts the reverse diffusion from Gaussian distribution at $t=T$ (see Fig. 2(a)), the total number of the reverse diffusion step to recover a clean images using CCDF can be significantly reduced. Furthermore, with better initialization, we prove that the number of reverse sampling can be further reduced as shown in Fig. 2(c). This implies that the existing neural-network (NN) based inverse solution can be synergistically combined with diffusion models to yield accurate and fast reconstruction by providing a better initial estimate.

Using extensive experiments across various problems such as super-resolution (SR), inpainting, and MRI reconstruction, we demonstrate that CCDF can significantly accelerate diffusion based models for inverse problems.

Background

For the given forward SDE in (1), there exists a reverse-time SDE running backwards :

where $dt$ is the infinitesimal negative time step, and $\bar{{\bm{w}}}$ is the Brownian motion running backwards.

Interestingly, one can train a neural network to approximate the actual score function via score matching to estimate $\bm{s}_{\theta}({\bm{x}},t)\simeq\nabla_{\bm{x}}\log p_{t}({\bm{x}})$, and plug it into (2) to numerically solve the reverse-SDE . Furthermore, to circumvent technical difficulties, de-noising score matching is typically used where $\nabla_{\bm{x}}\log p_{t}({\bm{x}})$ is replaced with $\nabla_{\bm{x}}\log p_{0t}({\bm{x}}(t)|{\bm{x}}(0))$.

2 Discrete Forms of SDEs

In this paper, we make use of two different SDEs: variance preserving (VP) SDE, and variance exploding (VE) SDE . First, by choosing

where $0<\beta(t)<1$ is a monotonically increasing function of noise scale, one achieves the variance preserving (VP)-SDE . On the other hand, variance exploding (VE) SDEs choose

where $\sigma(t)>0$ is again a monotonically increasing function, typically chosen to be a geometric series .

For the discrete diffusion models, we assume we have $N$ discretizations which are linearly distributed across $t\in[0,T]$. Then, VP-SDE can be seen as the continuous version of DDPM . Specifically, in DDPM, the forward diffusion is performed as

where $\bm{z}\sim\mathcal{N}(\textbf{0},\bm{I})$ and $\bar{\alpha}_{i}=\prod_{j=1}^{i-1}\alpha_{j}$ for $\alpha_{i}=1-\beta_{i}$ with monotonically increasing noise schedule $\beta_{1},\beta_{2},\dots,\beta_{N}\in(0,1)$. The associated reverse diffusion step is

where $\bm{s}_{\theta}({\bm{x}}_{i},i)$ is a discrete score function that matches $\nabla_{{\bm{x}}_{i}}\log p_{0i}({\bm{x}}_{i}|{\bm{x}}_{0})$. Further, the noise term $\sigma_{i}$ can be fixed to $\sigma_{i}=1-\alpha_{i}$ , or set to a learnable parameter .

For DDPM, denoising diffusion implicit model (DDIM) establishes the current state-of-the-art among the acceleration methods. Unlike DDPM, DDIM has no additive noise term during the reverse diffusion, allowing less iterations for competitive sample quality. Specifically, the reverse diffusion step is given as:

to express (2.2) as $\bar{\bm{x}}_{i-1}=\bar{\bm{x}}_{i}+(\sigma_{i-1}-\sigma_{i})\bm{z}_{\theta}({\bm{x}}_{i},i)$.

On the other hand, score matching with Langevin dynamic (SMLD) can be seen as the discrete version of VE-SDE. Specifically, the forward SMLD diffusion step is given by

where $\sigma_{i}=\sigma_{\text{min}}(\frac{\sigma_{\text{max}}}{\sigma_{\text{min}}})^{\frac{i-1}{N-1}}$, as defined in . The associated reverse diffusion is given by

where $\bm{z}\sim\mathcal{N}(\textbf{0},\bm{I})$.

Main Contribution

The goal of our CCDF acceleration scheme is to make the reverse diffusion start from $N^{\prime}:=Nt_{0}<N$ such that the resulting number of reverse diffusion step can be significantly reduced. For this, our CCDF algorithm is composed of two steps: forward diffusion up to $N^{\prime}$ with better initialization ${\bm{x}}_{0}$, which is followed by a reverse conditional diffusion down to $i=0$.

Specifically, for a given initial estimate ${\bm{x}}_{0}$, the forward diffusion process can be performed with a single step diffusion as follows:

where $\bm{z}\sim\mathcal{N}(\textbf{0},\bm{I})$, and $a_{N^{\prime}}$, $b_{N^{\prime}}$ for SMLD and DDPM can be computed for each diffusion model using (10) and (5), respectively.

In regard to the conditional difusion, SRDiff , SR3 are examples that are trained specifically for SR, with the low-resolution counterparts being encoded or concatenated as the input. However, these approaches attempt to redesign the score function so that one can sample from the conditional distribution, leading to a much complicated formulation.

Instead, here we propose a much simpler but effective conditional diffusion. Specifically, our reverse diffusion uses standard reverse diffusion, alternated with an operation to impose data consistency:

where the specific forms of $\bm{f}({\bm{x}}_{i},i)$ and $g({\bm{x}}_{i},i)$ depend on the type of diffusion models (see Table 1), $\bm{z}_{i}\sim\mathcal{N}(\bm{0},\bm{I})$, and $\bm{A}$ is a non-expansive mapping :

In particular, we assume $\bm{A}$ is linear. For example, one-iteration of the standard gradient descent or projection onto convex sets (POCS) in corresponds to our data consistency step in (14) with (15). See Supplementary Section D for algorithms used for each task.

2 Fast Convergence Principle of CCDF

Now, we are ready to show why CCDF provides much faster convergence than the standard conditional diffusion models that starts from Gaussian noise. In fact, the key innovation comes from the mathematical findings that while the forward diffusion increases the estimation error, the conditional reverse diffusion decreases it much faster at exponential rate. Accordingly, we can find a “sweet spot” $N^{\prime}$ such that the forward diffusion up to $N^{\prime}$ followed by reverse diffusion can significantly reduces the estimation error of the initial estimate ${\bm{x}}_{0}$. This fast convergence principle is shown in the following theorems, whose proofs can be found in Supplementary Materials. First, the following lemma is a simple consequence of independency of Gaussian noises.

Now, the following theorem, which is a key step of our proof, comes from the stochastic contraction property of the stochastic difference equation .

Consider the reverse diffusion using (13) and (14). Then, we have

Now we have the main results that shows the existence of the shortcut path for the acceleration.

For any $0<\mu\leq 1$, there exists a minimum $N^{\prime}(=t_{0}N<N)$ such that $\bar{\varepsilon}_{0,r}\leq\mu\varepsilon_{0}$. Furthermore, $N^{\prime}$ decreases as $\varepsilon_{0}$ gets smaller.

Theorem 1 states that the conditional reverse diffusion is exponentially contracting. Subsequently, Theorem 2 tells us that we can achieve superior results (i.e. tighter bound) with shorter sampling path. Hence, it is unnecessary for us to start sampling from $N$. Rather, we can start from an arbitrary timestep $N^{\prime}<N$, and still converge faster to the same point that could be achieved when starting the sampling procedure at $N$. Furthermore, as we have better initialization such that $\varepsilon_{0}$ is smaller, then we need smaller reverse diffusion step, achieving much higher acceleration.

For example, we can initialize the corrupted image with a pre-trained neural network $G_{\varphi}$, which has been widely studied across different tasks . These methods are typically extremely fast to compute, and thus does not introduce additional computational overload. Using this rather simple and fast fix, we observe that we are able to choose smaller values of $t_{0}$, endowed with much stabler performance. For example, in the case of MRI reconstruction, we can choose $t_{0}$ as small as 0.02, while outperforming score-MRI with 50$\times$ acceleration.

Experiments

We test our method on three different tasks: super-resolution, inpainting, and MRI reconstruction. For all methods, we evaluate the qualitative image quality and quantitative metrics as we accelerate the diffusion process by reducing the $t_{0}$ values. For the proposed method, we report on the results starting with neural network (NN)-initialized ${\bm{x}}_{0}$ unless specified otherwise.

Dataset. For vision tasks using face images, we use two datasets - FFHQ $256\times 256$, and AFHQ $256\times 256$. For FFHQ, we randomly select 50k images for training, and sample 1k images of test data separately. For AFHQ, we train our model using the images in the dog category, which consists of about 5k images. Testing was performed with the held-out validation set of 500 images of the same category. For the MRI reconstruction task, we use the fastMRI knee data, which consists of around 30k $320\times 320$-sized slices of coronal knee scans. Specifically, we use magnitude data given as the key reconstruction_esc. We randomly sample 10 volumes from the validation set for testing.

Quantitative metrics. Since it is well known that for high corruption factors, standard metrics such as PSNR/SSIM does not correlate well with the visual quality of the reconstruction , we report on the FID score based on pytorch-fidhttps://github.com/mseitzer/pytorch-fid. For MRI reconstruction, it is less sound to report on FID; hence, we report on PSNR.

Super-resolution. Experiments were performed across three different levels of SR factor - $\times 4,\times 8,\times 16$. We train a discretized VP-SDE based on IDDPM for each dataset - FFHQ and AFHQ, following the standards. Specific details can be found in Supplementary section D. For the one-step feed forward network corrector, we train the widely-used ESRGAN for each SR factor, using the same neural network architecture that was used to train the score function. We use three methods for comparison - ESRGAN, ILVR, and SR3https://github.com/Janspiry/Image-Super-Resolution-via-Iterative-Refinement. We note that the official code of SR3 is yet to be released, and hence we resort to unofficial re-implementation, which we train with default configurations. Additionally, in the original work of SR3 , the authors propose consecutively applying $\times 4$ SR models to achieve $16\times 16\mapsto 64\times 64\mapsto 256\times 256$ SR. In contrast, we report on a single $\times 16$ SR model which maps $16\times 16\mapsto 256\times 256$ directly.

Inpainting. The score function used in the inpainting task is the same model that was used to solve SR tasks, since we use task-agnostic conditional diffusion model. The feed-forward network was adopted from Yu et al. . We consider box-type inpainting with varying sizes: $96\times 96,128\times 128,160\times 160$. The model was trained for 50k steps with default configurations. We compare with score-SDE , using the same trained score function.

MRI reconstruction. Experiments were performed across three different levels of acceleration factor, with gaussian 1D sampling pattern - $\times 2,\times 4,\times 6$, each with 10%, 8%, 6% of the phase encoding lines included for autocalibrating signal (ACS) region. We train a VE-SDE based on ncsnpp, proposed in , and demonstrated specifically for MR reconstruction in . For comparison with compressed sensing (CS) strategy, we use total-variation (TV) regularized reconstruction. For feed forward network, we train a standard U-Net, using similar settings from . We use the same trained score function for comparison with score-MRI .

2 Super-resolution

Dependence on $\varepsilon_{0}$. We first demonstrate the dependency of stochastic contraction on the squared error term in Figure 3. For small squared difference, as in the case for many inverse problems, we see that the reverse diffusion stably converges to the same solution, even with small timestep $t_{0}$. In contrast, when random ${\bm{x}}_{0}$ is the starting point, $\varepsilon_{0}$ becomes large, and only with higher values of $t_{0}$ does the reverse SDE converge to a feasible solution.

Dependence on $t_{0}$. In Table 2, we report on the FID scores by varying the $t_{0}$ values with a fixed discretization step $\Delta t=1/1000$ in order to see which value is optimal for each degradation factor. Consistent with the theoretical findings, we see that as the corruption factor gets higher, and $\varepsilon_{0}$ gets larger, we typically need higher values of $t_{0}$ to achieve optimal results. Interesting enough, we observe that there always exist a value $t_{0}\in[0,1)$ where the FID score is lower (lower is better) than when using full reverse diffusion from $T=1$.

Comparison study. The results of various super-resolution algorithms is compared in Fig. 4. We compare with SR3 and ILVR , with setting the number of iterations for reconstruction same for ILVR, SR3, and the proposed method. We clearly see that SR3 and ILVR starting from pure Gaussian noise at $T=1$ cannot generate satisfactory results with 20 iterations, whereas our method can estimate high-fidelity samples with details preserved even with only 20 iterations starting from $t_{0}=0.2$. Visualizing the trend of FID score in Figure 5, we see that the quality of the image degrades as we use less and less number of iterations for the ILVR method, whereas the proposed method is able to keep the FID score at about the same level, or even boost the image quality, with less iterations.

We also perform a comparison study where we set the total number of diffusion steps to $N=1000$ starting from $T=1$ for ILVR , and set $t_{0}$ to 0.1, 0.2, and 0.3 for each factor, thereby reducing the number of diffusion steps to 100, 200, and 300, respectively, by our method. In Table 3, we demonstrate by using the proposed method, we achieve results that are on par or even better. For qualitative analysis, see Supplementary Section E.

Incorporation of DDIM. As briefly discussed before, CCDF can be combined together with approaches that searches for the optimal (full) reverse diffusion path. In Fig. 6, we illustrate that we can reduce the number of iterations to as little as 5 steps, and still maintain high image quality.

3 Inpainting

We illustrate the results of inpainting in Fig. 7. Consistent with what was observed in the SR task, the results in Figure 7 show that using full reverse diffusion with large discretization steps is inefficient, leading to unrealistic output. On the other hand, our method can reconstruct very realistic images within this small budget.

Comparison with prior arts by setting relatively large number of iterations is shown in Table 4. We observe that the proposed method outperforms both score-SDE with full reverse diffusion, and SN-PatchGAN, in terms of FID score. For detailed comparison and further experiments, see Supplementary Section E.

4 MRI reconstruction

We summarize and compare our results in Figure 8, and the quantitative metrics are presented in Table 5. In the task of MR reconstruction, we observe that we can push the $t_{0}$ value down to very small values: $t_{0}=0.02$, and still achieve remarkable results, even outperforming score-POCS which uses full reverse diffusion. When we compare the proposed method which uses 20 iterations vs. score-POCS with 20 iterations, we see that score-POCS cannot generate a feasible image, arriving at what looks like pure noise, as demonstrated in Figure 1. With other tasks, we could see that higher degradations typically require increased $t_{0}$ values. With CCDF, we do not see such trend, and observe that selecting low values of $t_{0}\in[0.02,0.1]$ stably gives good results. We emphasize that this is a huge leap towards practical usage of diffusion models in clinical settings, where fast reconstruction is crucial for real-time deployment.

Discussion

We note that we are not the first to propose starting from forward-diffused data in the context of diffusion models. It was first introduced in SDEdit , but in a different context with distinct aim form ours. In SDEdit, forward diffusion was used up to $t_{0}\in[0.3,0.6]$, which a relatively higher value than those used in our work $t_{0}\leq 0.2$, since the purpose was to destroy the signal so as to acquire high fidelity images from coarse strokes.

Our work differs from SDEdit in that we consider this procedure in a more rigorous framework and first reveal that starting from a better initialization for inverse problems significantly accelerate the reverse diffusion. This leads to a novel hybridization that has not been covered before: a simple incorporation of pre-trained feed-forward NNs can be very efficient at pushing $t_{0}$ to smaller limits, even as small as $t_{0}=0.02$ in the case of MRI reconstruction.

We note that the choice of $t_{0}$ for acceleration varies by quite a margin across different tasks, and the degree of corruptions. Currently, there does not exist clear and concise rules for selecting such values as we do not have a knowledge of $\varepsilon_{0}$ a priori. Thus, one needs to rely mostly on trial-and-error, which could potentially reduce practicality. Building an adaptive method that can automatically search for the optimal $t_{0}$ values will be beneficial, and we leave this venue for possible direction of future research.

Conclusion

In this work, we proposed a method to accelerate conditional diffusion models, by studying the property of stochastic contraction. When solving inverse problems via conditional reverse diffusion, rather than starting at random Gaussian noise, we proposed to initialize the starting from forward-diffused data from a better initialization, such as one-step correction via NN. Using the stochastic contraction theory, we showed theoretically why taking the shortcut path is in fact optimal, and back our statement by showing diverse applications in which we both achieve acceleration along with increased stability and performance.

References

Supplementary Material

Appendix A Mathematical Preliminaries

Using the intermediate value theorem for the function $\bm{f}({\bm{x}})$, we can easily see that $\bm{f}({\bm{x}})$ is contracting with the rate $0\leq\lambda<1$ if $\bm{f}$ satisfies the following:

Now, we provide a theorem for discrete stochastic contraction, which is slightly modified from the contraction theorem of stochastic difference equation in .

Consider the stochastic difference equation:

The following corollary is a simple consequence of Theorem A.1.

Consider the stochastic difference equation associated with the data fidelity term:

as $\sigma_{\max}(\bm{A})\leq 1$ for a non-expansive linear mapping. Furthermore, we have

Let $s_{\theta}({\bm{x}}_{i},i)$ be a sufficiently expressive parameterized score function so that

where $\bm{z}\sim\mathcal{N}(0,\bm{I})$ and $(a_{i},b_{i})$ are defined in (5) and (10) for DDPM and SMLD, respectively. Using (28), we have

where T denotes the transpose. This concludes the proof. ∎

Appendix B Proof of Theorem 1

Let $N$ be the standard reverse diffusion step when starting from $T=1$. Then, the number of discretization step for our method is given $N^{\prime}=Nt_{0}<N$ so that $t_{0}$ can refer to the acceleration factor. We further define a new index $i=N^{\prime}-j$ to convert the reverse diffusion index $j=N^{\prime},\cdots,1$ to a forward direction index $i=0,1,\cdots,N^{\prime}$. This does not change the contraction property of the stochastic difference equation. Therefore, without loss of generality, we use the aforementioned contraction property of stochastic difference equation for the index $i=0,1,\cdots,N^{\prime}$. Now, we are ready to provide the proof.

In DDPM, the discrete version of the forward diffusion is given by Eq. (5), and the reverse diffusion is given by eq. (6). Here, $\bm{z}_{\theta}({\bm{x}},i)$ is trained by

It was shown that $\bm{z}_{\theta}({\bm{x}},i)$ is a scaled version of the score function :

Therefore, the contraction rate is given by

as $0<\alpha_{i},\bar{\alpha}_{i}<1$. Furthermore, we can easily show that

as $\bar{\alpha}_{i}$ is decreasing with $i$.

B.2 SMLD: Discrete Version of VE-SDE

In discrete version of VE-SDE, the forward diffusion is given by (10). The associated reverse diffusion is given by (11). Thus, we have

as $\sigma_{i}$ is increasing with $i$. Furthermore, we can easily show that

B.3 DDIM

The DDIM forward diffusion can be set identically to the forward diffusion of DDPM (5), whereas the reverse diffusion is given as (2.2). In fact, with a proper reparameterization, one can cast DDIM such that it is equivalent to the discrete version of VE-SDE without noise terms. More specifically, if we define the following reparametrization:

Furthermore, the corresponding score function with respect to the reparameterization is

The forward diffusion (5) can be equivalently represented by the reparameterization as:

as $\sigma_{i}$ is increasing with $i$. Furthermore, we can easily show that $C=0$ as there is no noise term.

Appendix C Proof of Theorem 2

For some of the proofs, we borrow more tight inequality to obtain the result. In fact, the inequality of stochastic contraction

is a rough estimation of recursive inequality

where $\bar{\varepsilon}_{j,r}$ denotes the estimation error between reverse conditional diffusion path down to $j$. Accordingly, we have

which is reduced to (50) when $\lambda_{j}$ and $C_{j}$ are uniformly bounded by $\lambda$ and $C$, respectively.

Now, our proof strategy is as follows. We specify reasonable conditions on $\{\beta_{i}\}$ or $\{\sigma^{2}_{i}\}$, which are satisfied by the existing DDPM, SLMD, and DDIM scheduling approaches. Then, for any $0<\mu\leq 1$, our goal is to to show that there exists $N^{\prime}$ such that

and $N^{\prime}$ decreases as $\varepsilon_{0}$ gets smaller.

Without loss of generality, we assume that ground truth image and the corrupted image are normalized within range $$, i.e.${\bm{x}},\bar{\bm{x}}\in^{n}$. Then, we have

We separately investigate each term in (52). First, from theorem 1,

where the last inequality comes from (53). Subsequently,

where the first inequality comes from $\prod_{i=1}^{j-1}\lambda_{i}^{2}\leq\prod_{i=1}^{j-1}1\leq 1$ and the last equality is from (55). Therefore,

where the third inequality holds by (54), and the inequality in (56) comes from Lemma C.1 (see below). Furthermore, from (55), we can see that $N^{\prime}$ becomes smaller for a smaller $\varepsilon_{0}$. This concludes the proof of DDPM.

where the first inequality comes from $\bar{\alpha}_{j}=\bar{\alpha}_{j-1}\alpha_{j}\leq\bar{\alpha}_{j-1}$, and the second inequality is the inequality of arithmetic and geometric means, and the third equality is from the linear increasing $\beta_{j}$ from $\beta_{0}=0$. Finally, using

C.2 SMLD

Assume that the minimum and maximum values of variance satisfy the following:

Now, we can choose $N^{\prime}$ such that it satisfies the following conditions:

On the other hand, in the geometric scheduling of noise, for all $i$, we have

Hence, by plugging in (66) to (50), we have

where the third inequality comes from the bounds in (67), (68), and the fact that $\tau=\frac{tr(A^{T}A)}{n}<1$ for a non-expansive linear mapping $A$.

Finally, we can easily see that the value $N^{\prime}$ satisfying (63) decreases as $\varepsilon_{0}$ decreases.

C.3 DDIM

In DDIM, we have $C_{j}=0$ for Eq. (52). Let $\sigma_{0}$ and $N^{\prime}$ satisfy the following:

where the second equality comes from $\lambda_{j}=\sigma_{j-1}/\sigma_{j}$ and the last equality comes from Eqs. (69) and (70).

We can also easily see that the minimum value $N^{\prime}$ satisfying (70) decreases as $\varepsilon_{0}$ decreases, as $\sigma_{i}^{2}$ is an increasing sequence in DDIM.

Appendix D Implementation detail

In this section, we provide detailed explanation of discrete version of CCDF for each application. Again, the number of discretization step for our method is given $N^{\prime}=Nt_{0}<N$ where $t_{0}$ refers to the acceleration factor.

For these problems, we employ the discretized version of the VP-SDE, which has shown impressive results on conditional generation . Namely, we use DDPM , with several strategies introduced in improved DDPM (IDDPM) for both training the score function and for reverse diffusion procedure.

The modified reverse diffusion is given by

Specifically, $v$ and the score function $\bm{s}_{\theta}$ are trained using the following objective

where $L_{simple}(\theta)$ is given in (37) and we apply stop-gradient for the $L_{VLB}$ so that the gradient of the loss contributes only to estimating the model variance.

For the training of score function, we use a U-Net architecture as used in with the loss function as given in (73). Multi-headed attention was used only at the 16$\times$16 resolution. Linear beta noise scheduling with $\beta_{\text{min}}=0.0001$ and $\beta_{\text{max}}=0.02$ were used, with $N=1000$ discretization. We train the model with a batch size of 2, and a static learning rate of 1e-4 with Adam optimizer for 5M steps. Exponential moving average (EMA) rate of 0.9999 was applied to the model.

For super-resolution, we define a blur kernel ${\bm{h}}_{D}$ which is defined by successive applications of the downsampling filter by a factor $D$, and upsampling filter by a factor $D$. This can be represented as a matrix multiplication:

where ${\bm{x}}^{\prime}$ denotes intermediate estimate from the reverse diffusion. Then, we use the following data consistency iteration:

where ${\bm{x}}_{i}$ is the current estimate, and $\hat{\bm{x}}_{i}$ is the forward propagated image from the initial measurement $\hat{\bm{x}}(0)$:

We can easily see that $\sigma_{\max}({\bm{A}})\leq 1$ for the normalized filter ${\bm{h}}_{D}$.

Similarly, for the case of image inpainting, $\bm{P}$ is just a diagonal matrix with 1 at the measured locations and 0 on the unmeasured locations so that $\sigma_{\max}({\bm{A}})\leq 1$.

The resulting pseudo-code implementation of the algorithm is given in Algorithm 1.

D.2 DDIM for Super-resolution/Inpainting

Note that we can use the same score function trained for DDPM, and use it in DDIM sampling . Here, we study the effect on combining DDIM together with the proposed method to achieve even further acceleration. All we need to do is modify the unconditional update step, arriving at Algorithm 2.

D.3 MRI reconstruction

For the task of MRI reconstruction, we ground our work on Score-MRI , and modify the previous algorithm for our purpose. The algorithm is given in Algorithm 3. Specifically, we use variance exploding (VE-SDE) with predictor-corrector (PC) sampling which gives optimal results for MR reconstruction. For the step size of the corrector (Langevin dynamics) step, we use the following

with $r=0.16$ set as constant. For training the score function, we use the following minimization strategy:

with $\sigma_{\min}=0.01,\sigma_{max}=378$. We construct a modified U-Net model introduced in , namely ncsnpp. Adam optimizer is used for optimization, with a static learning rate of 2e-4 for 5M steps. EMA rate of 0.999 is used, and gradient clipping is applied with the maximum value of 1.0.

In compressed sensing MRI, the subsampled $k$-space data ${\bm{y}}$ is obtained from underlying image ${\bm{x}}$ as:

where $\bm{F}$ denote the Fourier transform and its inverse, $\bm{D}$ is a diagonal matrix indicating the $k$-space sampling location and ${\bm{y}}$ is the original zero-filled $k$-space data.

The associated data consistency imposing operator is then defined by

where $\bm{F}^{-1}$ is the inverse Fourier transform. Therefore, we have

Again, we can easily see $\sigma_{\max}({\bm{A}})\leq 1$ as the Fourier transform is orthonormal.

The corresponding pseudo-code implementation is shown in Algorithm 3.

All training and inference algorithms were implemented in PyTorch, and were performed on a single RTX 3090 GPU.

Appendix E Additional Experiments

Comparison study. In Fig. E.1, we compare the results of super-resolution using fairly large number of diffusion steps, as opposed to using only 20 number of diffusion steps as shown in Fig. 4. This is a region where ILVR is known to perform well, as opposed to the few-step setting. While in Fig. E.1, ILVR uses 1000 steps of diffusion, the proposed method only uses 100, 200, and 300 steps of diffusion for $\times 4,\times 8,$ and $\times 16$, respectively. Nevertheless, the quality of reconstruction does not degrade, thanks to the contraction property of CCDF.

Incorporation of DDIM. We provide additional SR results using CCDF + DDIM. In Figure E.2, we show an experiment with the FFHQ dataset, where we compare the combination of ILVR + DDIM, and proposed method + DDIM. For ILVR + DDIM, in order to reduce the number of iterations, we choose larger discretization steps used in DDIM. For the proposed method, we fix $N=50$, and reduce the value of $t_{0}$ to achieve less iterations. In the figure, we confirm that our method can be used together with DDIM to create high-fidelity samples with as small as 5 reverse diffusion iterations, even when it comes down to extreme cases of SR $\times 8$ or SR $\times 16$. Additionally, we observe that the results with $t_{0}\leq 0.5$ is superior to the $t_{0}=1.0$ counterparts, again confirming our theory.

The same trend can also be seen via quantitative metrics in Table E.1. Using limited number of diffusion steps, FID score in the case of ILVR+DDIM grows exponentially, as we decrease the number of steps taken. Contrarily, our method is able to improve the metric by quite a margin, as opposed to using full diffusion with 50 steps in total. This trend is indeed similar to the experiments performed with DDPM.

Experiments with ImageNet. ImageNet contains diverse categories of natural images, and are known to be much harder to model, due to its highly multimodal nature. We try to examine if CCDF scales even to this challenging task, using a pre-trained model provided in the guided-diffusion github repositoryhttps://github.com/openai/guided-diffusion. As with other experiments with FFHQ or AFHQ dataset, we train an ESRGAN model for each SR factor, and use it as our initialization strategy. In Fig. F.2, we can see that our CCDF strategy outperforms ILVR using full reverse diffusion, and also vastly improves the image quality of ESRGAN, which is our initialization.

E.2 Inpainting

Dependence on $t_{0}$. As in Table 2, we compare the FID score of reconstructions for the inpainting task, as we vary the $t_{0}$ values in table E.2. We notice similar results from the SR task, in the sense that there always exist $t_{0}\in(0,T)$ which gives higher scores than using full diffusion. With relatively small boxes, we see that $t_{0}=0.1$ is optimal, whereas we typically need more diffusion steps for larger boxes.

Comparison study. We compare the proposed CCDF strategy with SN-PatchGAN , and score-SDE using 1000 steps in Fig. E.3. For SN-PatchGAN in Fig. E.3 (b), we often see highly unrealistic details e.g. near the mouth. Note that SN-patchGAN serves as the initialization point for CCDF in inpainting. Leveraging this imperfect initialization, the proposed method is able to provide reconstructions that are highly realistic, as can be seen in fig. E.3 (d). It is also notable that score-SDE using full diffusion more often than not produces results that are incoherent with the known regions (see second row of Fig. E.3 (c)), while the proposed method stably outputs coherent results.

Ablation study. We perform an ablation study comparing the effect of different initialization strategy. Table E.3 shows the difference in the results when using vanilla initialization with the corrupted image, and NN initialization. We see that with all $t_{0}$, NN initialization performs marginally better than vanilla initialization. The difference becomes clearer as we decrease the value of $t_{0}$ to 0.1. The same ablation study was performed also for MRI reconstruction task, and is illustrated in Figure E.4. We see similar trend as in the SR task.

Furthermore, we provide additional qualitative results of each task on various datasets, focusing mainly on showing the trend of reconstruction results as we vary the value of $t_{0}$. In Figure F.3, we compare the achievable image quality by fixing the number of reverse diffusion steps to 20. Consistent with what we saw in Figure 4, we see that our method largely outperforms the other diffusion model-based methods. In Figure F.4, Figure F.5 respectively, we see that we can stably arrive at a feasible solution with different values of $t_{0}\in[0.1,0.5]$, typically requiring higher values of $t_{0}$ for severer degradation.

Appendix F Validity of assumption

In Lemma A.1, we assumed that $s_{\theta}({\bm{x}}_{i},t)=\frac{\partial}{\partial{\bm{x}}_{i}}\log p_{0i}({\bm{x}}_{i}|{\bm{x}}_{0})$. In this section, we briefly show that the assumption is valid. In the theoretic side, showed that given an optimal reconstruction function $r^{*}_{\sigma}({\bm{x}}_{t})$ trained with denoising autoencoder loss asymptotically behaves as

where the equation emphasizes the behavior of the optimal score function, regarding how it was parameterized. This asymptotic behavior hints that the error will be small, especially when $\sigma$ approaches zero. In our case, $\sigma\rightarrow 0$ as $t\rightarrow 0$, and hence the behavior holds near $t=0$.

In order to numerically validate such error, we conducted an experiment to calculate the actual average error norm, illustrated in Fig. F.1. Here, we see that the error norm mostly stays at very low values across the range. We do observe that the magnitude of error inevitably increases when the noise level is too large, so the error grows where $t>0.5$. Nevertheless, we note that most of our contraction analysis stays in the $t\in[\epsilon,0.5]$ regime, and hence the assumption made in Lemma A.1. is practical.