Denoising Diffusion Restoration Models
Abstract
Many interesting tasks in image restoration can be cast as linear inverse problems.
A recent family of approaches for solving these problems uses stochastic algorithms that sample from the posterior distribution of natural images given the measurements.
However, efficient solutions often require problem-specific supervised training to model the posterior, whereas unsupervised methods that are not problem-specific typically rely on inefficient iterative methods.
This work addresses these issues by introducing Denoising Diffusion Restoration Models (DDRM), an efficient, unsupervised posterior sampling method.
Motivated by variational inference, DDRM takes advantage of a pre-trained denoising diffusion generative model for solving any linear inverse problem.
We demonstrate DDRM’s versatility on several image datasets for super-resolution, deblurring, inpainting, and colorization under various amounts of measurement noise.
DDRM outperforms the current leading unsupervised methods on the diverse ImageNet dataset in reconstruction quality, perceptual quality, and runtime, being 5 × faster than the nearest competitor.
DDRM also generalizes well for natural images out of the distribution of the observed ImageNet training set.
Introduction
Many problems in image processing, including super-resolution [31, 17], deblurring [28, 48], inpainting [55], colorization [29, 58], and compressive sensing [1], are instances of linear inverse problems, where the goal is to recover an image from potentially noisy measurements given through a known linear degradation model.
For a specific degradation model, image restoration can be addressed through end-to-end supervised training of neural networks using pairs of original and degraded images [14, 58, 41].
However, real-world applications such as medical imaging often require flexibility to cope with multiple, possibly infinite, degradation models [46].
Here, unsupervised approaches based on learned priors [36], where the degradation model is only known and used during inference, may be more desirable since they can adapt to the given problem without re-training [51].
By learning sound assumptions over the underlying structure of images ( e.g. , priors, proximal operators or denoisers), unsupervised approaches can achieve effective restoration without training on specific degradation models [51, 40].
Under this unsupervised setting, priors based on deep neural networks have demonstrated impressive empirical results in various image restoration tasks [40, 50, 43, 38, 15].
To recover the signal, [Figure] most existing methods obtain a prior-related term over the signal from a neural network ( e.g. , the distribution of natural images), and a likelihood term from the degradation model.
They combine the two terms to form a posterior over the signal, and the inverse problem can be posed as solving an optimization problem ( e.g. , maximum a posteriori [8, 40]) or solving a sampling problem ( e.g. , posterior sampling [2, 3, 25]).
Then, these problems are often solved with iterative methods, such as gradient descent or Langevin dynamics, which may be demanding in computation and sensitive to hyperparameter tuning.
An extreme example is found in [30] where a “fast” version of the algorithm uses 15 , 000 neural function evaluations (NFEs).
Inspired by this unsupervised line of work, we introduce an efficient approach named Denoising
Diffusion Restoration Models (DDRM), that can achieve competitive results in as low as 20 NFEs.
DDRM is a denoising diffusion generative model [44, 19, 45] that gradually and stochastically denoises a sample to the desired output, conditioned on the measurements and the inverse problem.
This way we introduce a variational inference objective for learning the posterior distribution of the inverse problem at hand.
We then show its equivalence to the objective of an unconditional denoising diffusion generative model [19], which enables us to deploy such models in DDRM for various linear inverse problems (see Figure 2).
To our best knowledge, DDRM is the first general sampling-based inverse problem solver that can efficiently produce a range of high-quality, diverse, yet valid solutions for general content images.
We demonstrate the empirical effectiveness of DDRM by comparing with various competitive methods based on learned priors, such as Deep Generative Prior (DGP) [38], SNIPS [25], and Regularization by Denoising (RED) [40].
On ImageNet examples, DDRM mostly outperforms the neural network baselines under noiseless super-resolution and deblurring measured in PSNR and KID [5], and is at least 50 × more efficient in terms of NFEs when it is second-best.
Our advantage becomes even larger when measurement noise is involved, as noisy artifacts produced by iterative methods do not appear in our case.
Over various real-world images, we further show DDRM results on super-resolution, deblurring, inpainting and colorization (see Figure 1).
A DDRM trained on ImageNet also works on images that are out of its training set distribution (see Figure 6).
Background
Linear Inverse Problems.
A general linear inverse problem is posed as
[Equation]
where we aim to recover the signal x ∈ R n from measurements y ∈ R m , where H ∈ R m × n is a known linear degradation matrix, and z ∼ N (0 , σ 2 y I ) is an i.i.d.
additive Gaussian noise with known variance.
The underlying structure of x can be represented via a generative model, denoted [Figure] as p θ ( x ) .
Given y and H , a posterior over the signal can be posed as: p θ ( x | y )
∝ p θ ( x ) p ( y | x ) , where the “likelihood” term p ( y | x ) is defined via Equation (1); such an approach leverages a learned prior p θ ( x ) , and we call it an “unsupervised” approach based on the terminology in [36], as the prior does not necessarily depend on the inverse problem.
Recovering x can be done by sampling from this posterior [2], which may require many iterations to produce a good sample.
Alternatively, one can also approximate this posterior by learning a model via amortized inference ( i.e. , supervised learning); the model learns to predict x given y , generated from x and a specific H .
While this can be more efficient than sampling-based methods, it may generalize poorly to inverse problems that have not been trained on.
Denoising Diffusion Probabilistic Models.
Structures learned by generative models have been applied to various inverse problems and often outperform data-independent structural constraints such as sparsity [7].
These generative models learn a model distribution p θ ( x ) that approximates a data distribution q ( x ) from samples.
In particular, diffusion models have demonstrated impressive unconditional generative modeling performance on images [13].
Diffusion models are generative models with a Markov chain structure x T → x T − 1 → . . .
→ x 1 → x 0 (where x t ∈ R n ), which has the following joint distribution:
[Equation]
After drawing x 0: T , only x 0 is kept as the sample of the generative model.
To train a diffusion model, a fixed, factorized variational inference distribution is introduced:
[Equation]
which leads to an evidence lower bound (ELBO) on the maximum likelihood objective [44].
A special property of some diffusion models is that both p ( t ) θ and q ( t ) are chosen as conditional Gaussian distributions for all t < T , and that q ( x t | x 0 ) is also a Gaussian with known mean and covariance, i.e. , x t can be treated as x 0 directly corrupted with Gaussian noise.
Thus, the ELBO objective can be reduced into the following denoising autoencoder objective (please refer to [45] for derivations):
[Equation]
where f ( t ) θ is a θ -parameterized neural network that aims to recover a noiseless observation from a noisy x t , and γ 1: T are a set of positive coefficients that depend on q ( x 1: T | x 0 ) .
Denoising Diffusion Restoration Models
Inverse problem solvers based on posterior sampling often face a dilemma: unsupervised approaches apply to general problems but are inefficient, whereas supervised ones are efficient but can only address specific problems.
To solve this dilemma, we introduce Denoising Diffusion Restoration Models (DDRM), an unsupervised solver for general linear inverse problems, capable of handling such tasks with or without noise in the measurements.
DDRM is efficient and exhibits competitive performance compared to popular unsupervised solvers [40, 38, 25].
The key idea behind DDRM is to find an unsupervised solution that also suits supervised learning objectives.
First, we describe the variational objective for DDRM over a specific inverse problem (Section 3.1).
Next, we introduce specific forms of DDRM that are suitable for linear inverse problems and allow pre-trained unconditional and class-conditional diffusion models to be used directly (Sections 3.2, 3.3).
Finally, we discuss practical algorithms that are compute and memory efficient (Sections 3.4, 3.5).
3.1 Variational Objective for DDRM
For any linear inverse problem, we define DDRM as a Markov chain x T → x T − 1 → . . .
→ x 1 → xconditioned on y , where
[Equation]
and x 0 is the final diffusion output.
In order to perform inference, we consider the following factorized variational distribution conditioned on y :
[Equation]
leading to an ELBO objective for diffusion models conditioned on y (details in Appendix A).
In the remainder of the section, we construct suitable variational problems given H and σ y and connect them to unconditional diffusion generative models.
To simplify notations, we will construct the variational distribution q such that q ( x t | x 0 ) =
N ( x 0 , σ 2 t I ) for noise levels
= σ 0 < σ 1 < σ 2 < . . .
<
σ T .
In Appendix B, we will show that this is equivalent to the distribution introduced in
DDPM [19] and DDIM [45], 3 up to fixed linear transformations over x t .
3.2 A Diffusion Process for Image Restoration
Similar to SNIPS [25], we consider the singular value decomposition (SVD) of H , and perform the diffusion in its spectral space.
The idea behind this is to tie the noise present in the measurements y with the diffusion noise in x 1: T , ensuring that the diffusion result x 0 is faithful to the measurements.
By using the SVD, we identify the data from x that is missing in y , and synthesize it using a diffusion process.
In conjunction, the noisy data in y undergoes a denoising process.
For example, in inpainting with noise ( e.g. , H = diag([1 , . .
, 0]) , σ y ≥ 0 ), the spectral space is simply the pixel space, so the model should generate the missing pixels and denoise the observed ones in y .
For a general linear H , its SVD is given as
[Equation]
where U ∈ R m × m , V ∈ R n × n are orthogonal matrices, and Σ ∈ R m × n is a rectangular diagonal matrix containing the singular values of H , ordered descendingly.
As this is the case in most useful degradation models, we assume m ≤ n , but our method would work for m > n as well.
We denote the singular values as s 1 ≥ s 2 ≥ . .
≥ s m , and define s i
=
for i ∈ [ m + 1 , n ] .
We use the shorthand notations for values in the spectral space: ¯ x ( i ) t is the i -th index of the vector ¯ x t =
V ⊤ x t , and ¯ y ( i ) is the i -th index of the vector
¯ y
= Σ †
U ⊤ y (where † denotes the Moore–Penrose pseudo-inverse).
Because V is an orthogonal matrix, we can recover x t from ¯ x t exactly by left multiplying V .
For each index i in ¯ x t , we define the variational distribution as:
[Equation]
where η ∈ (0 , 1] is a hyperparameter controlling the variance of the transitions, and η and η b may depend on σ t , s i , σ y .
We further assume that σ T ≥ σ y /s
i for all positive s i .
In the following statement, we show that this construction has the “Gaussian marginals” property similar to the inference distribution used in unconditional diffusion models [19].
Proposition 3.1.
The conditional distributions q ( t ) defined in Equations 4 and 5 satisfy the following:
[Equation]
defined by marginalizing over x t ′ (for all t ′ > t ) and y , where q ( y | x 0 ) is defined as in Equation (1) with x = x 0 .
We place the proof in Appendix C. Intuitively, our construction considers different cases for each index of the spectral space.
( i )
If the corresponding singular value is zero, then y does not directly provide any information to that index, and the update is similar to regular unconditional generation.
( ii )
If the singular value is non-zero, then the updates consider the information provided by y , which further depends on whether the measurements’ noise level in the spectral space ( σ y /s i ) is larger than the noise level in the diffusion model ( σ t ) or not; the measurements in the spectral space ¯ y ( i ) are then scaled differently for these two cases in order to ensure Proposition 3.1 holds.
Now that we have defined q ( t ) as a series of Gaussian conditionals, we define our model distribution p θ as a series of Gaussian conditionals as well.
Similar to DDPM, we aim to obtain predictions of xat every step t ; and to simplify notations, we use the symbol x θ,t to represent this prediction made by a model 5 f θ ( x t +1 , t + 1) :
R n × R → R n that takes in the sample x t +1 and the conditioned time step ( t + 1) .
We also define ¯ x ( i ) θ,t as the i -th index of ¯ x θ,t = V ⊤ x θ,t .
We define DDRM with trainable parameters θ as follows:
[Equation]
Compared to q ( t ) in Equations (4) and (5), our definition of p ( t ) θ merely replaces ¯ x ( i )
(which we do not know at sampling) with ¯ x ( i ) θ,t (which depends on our predicted x θ,t ) when t < T , and replaces ¯ x ( i ) 0 with
when t =
T .
It is possible to learn the variances [35] or consider alternative constructions where Proposition 3.1 holds; we leave these options as future work.
3.3 “Learning” Image Restoration Models
Once we have defined p ( t ) θ and q ( t ) by choosing σ 1: T , η and η b , we can learn model parameters θ by maximizing the resulting ELBO objective (in Appendix A).
However, this approach is not desirable
O r i g i n a l
D e g r a d e d
DD R M ( 20 ) (a) Inpainting results on cat images.
O r i g i n a l
D e g r a d e d
DD R M ( 20 ) (b) Deblurring results ( σ y = 0 . 05 ) on bedroom images.
since we have to learn a different model for each inverse problem (given H and σ y ), which is not flexible enough for arbitrary inverse problems.
Fortunately, this does not have to be the case.
In the following statement, we show that an optimal solution to DDPM / DDIM can also be an optimal solution to a DDRM problem, under reasonable assumptions used in prior work [19, 45].
Theorem 3.2.
Assume that the models f ( t ) θ and f ( t ′ ) θ do not have weight sharing whenever t ̸ = t ′ , then when η = 1 and η b = 2 σ 2 t σ 2 t + σ 2 y /s 2 i , the ELBO objective of DDRM (details in Appendix A) can be rewritten in the form of the DDPM / DDIM objective in Equation (2) .
We place the proof in Appendix C. Even for different choices of η and η b , the proof shows that the DDRM objective is a weighted sumof-squares error in the spectral space, and thus pre-trained DDPM models are good approximations to the optimal solution.
Therefore, we can apply the same diffusion model ( unconditioned on the inverse problem) using the updates in Equation (7) and Equation (8) and only modify H and its SVD ( U , Σ , V ) for various linear inverse problems.
3.4 Accelerated Algorithms for DDRM
Typical diffusion models are trained with many timesteps ( e.g. , 1000) to achieve optimal unconditional image synthesis quality, but sampling speed is slow as many NFEs are required.
Previous works [45, 13] have accelerated this process by “skipping” steps with appropriate update rules.
This is also true for DDRM, since we can obtain the denoising autoencoder objective in Equation (2) for any choice of increasing σ 1: T .
For a pre-trained diffusion model with T ′ timesteps, we can choose σ 1: T to be a subset of the T ′ steps used in training.
3.5 Memory Efficient SVD
Our method, similar to SNIPS [25], utilizes the SVD of the degradation operator H .
This constitutes a memory consumption bottleneck in both algorithms as well as other methods such as Plug and
Play (PnP) [51], as storing the matrix V has a space complexity of Θ( n 2 ) for signals of size n .
By leveraging special properties of the matrices H used, we can reduce this complexity to Θ( n ) for denoising, inpainting, super resolution, deblurring, and colorization (details in Appendix D).
Related Work
Various deep learning solutions have been suggested for solving inverse problems under different settings (see a detailed survey in [37]).
We focus on the unsupervised setting, where we have access to a dataset of clean images at training time, but the degradation model is known only at inference time.
This setup is inherently general to all linear inverse problems, a property desired in many real-world applications such as medical imaging [46, 20]. [Table]
Almost all unsupervised inverse problem solvers utilize a trained neural network in an iterative scheme.
PnP, RED, and their successors [51, 40, 32, 49] apply a denoiser as part of an iterative optimization algorithm such as steepest descent, fixed-point, or alternating direction method of multipliers (ADMM).
OneNet [39] trained a network to directly learn the proximal operator of
ADMM.
A similar use of denoisers in different iterative algorithms is proposed in [34, 16, 30].
The authors of [43] leverages robust classifiers learned with additional class labels.
Another approach is to search the latent space of a generative model for a generated image that, when degraded, is as close as possible to the given measurements.
Multiple such methods were suggested, mainly focusing on generative adversarial networks (GANs) [7, 11, 33].
While they exhibit impressive results on images of a specific class, most notably face images, these methods are not shown to be largely successful under a more diverse dataset such as ImageNet [12].
Deep Generative Prior (DGP) mitigates this issue by optimizing the latent input as well as the weights of the GAN’s generator [38].
More recently, denoising diffusion models were used to solve inverse problems in both supervised ( i.e. , degradation model is known during training)
[42, 41, 13, 10, 54] and unsupervised settings [22, 26, 25, 21, 46, 47, 9].
Unlike previous approaches, most diffusion-based methods can successfully recover images from measurements with significant noise.
However, these methods are very slow, often requiring hundreds or thousands of iterations, and are yet to be proven on diverse datasets.
Our method, motivated by variational inference, obtains problem-specific, non-equilibrium update rules that lead to high-quality solutions in much fewer iterations.
ILVR [9] suggests a diffusion-based method that handles noiseless super-resolution, and can run in 250 steps.
In Appendix H, we prove that when applied on the same underlying generative diffusion model, ILVR is a special case of DDRM.
Therefore, ILVR can be further accelerated to run insteps, but unlike DDRM, it provides no clear way of handling noise in the measurements.
Similarly, the authors of [22] suggest a score-based solver for inverse problems that can converge in a small number of iterations, but does not handle noise in the measurements.
Experiments 5.1 Experimental Setup
We demonstrate our algorithm’s capabilities using the diffusion models from [19], which are trained on CelebA-HQ [23], LSUN bedrooms, and LSUN cats [56] (all 256 × 256 pixels).
We test these models on images from FFHQ [24], and pictures from the internet of the considered LSUN category, respectively.
In addition, we use the models from [13], trained on the training set of ImageNet 256 × 256 and 512 × 512 , and tested on the corresponding validation set.
Some of the ImageNet models require class information.
For these models, we use the ground truth labels as input, and denote our algorithm as DDRM class conditional (DDRM-CC).
In all experiments, we use η = 0 . 85 , η b = 1 , and a uniformly-spaced timestep schedule based on the 1000-step pre-trained models (more details in Appendix E).
The number of NFEs (timesteps) is reported in each experiment.
In each of the inverse problems we show, pixel values are in the range [0 , 1] , and the degraded measurements are obtained as follows: ( i ) for super-resolution, we use a block averaging filter to downscale the images by a factor of 2 , 4 , or 8 in each axis; ( ii ) for deblurring, the images are blurred [Table] by a 9 × 9 uniform kernel, and singular values below a certain threshold are zeroed, making the problem more ill-posed.
( iii ) for colorization, the grayscale image is an average of the red, green, and blue channels of the original image; ( iv ) and for inpainting, we mask parts of the original image with text overlay or randomly drop 50% of the pixels.
Additive white Gaussian noise can optionally be added to the measurements in all inverse problems.
We additionally conduct experiments on bicubic super-resolution and deblurring with an anisotropic Gaussian kernel in Appendix I.
Our code is available at https://github.com/bahjat-kawar/ddrm .
5.2 Quantitative Experiments
In order to quantify DDRM’s performance, we focus on the ImageNet dataset ( 256 × 256 ) for its diversity.
For each experiment, we report the average peak signal-to-noise ratio (PSNR) and structural similarity index measure (SSIM) [52] to measure faithfulness to the original image, and the kernel
Inception distance (KID) [5], multiplied by 10 3 , to measure the resulting image quality.
We compare DDRM (with 20 and 100 steps) with other unsupervised methods that work in reasonable time (requiring 1500 NFEs or less) and can operate on ImageNet. Namely, we compare with RED [Figure] O r i g i n a l
B l u rr e d DD R M ( 20 ) [40], DGP [38], and SNIPS [25].
The exact setup of each method is detailed in Appendix F.
We used the same hyperparameters for noisy and noiseless versions of the same problem for DGP, RED, and
SNIPS, as tuning them for each version would compromise their unsupervised nature.
Nevertheless, the performance of baselines like RED with such a tuning does not surpass that of DDRM, as we show in Appendix F.
In addition, we show upscaling by bicubic interpolation as a baseline for super-resolution, and the blurry image itself as a baseline for deblurring.
OneNet [39] is not included in the comparisons as it is limited to images of size 64 × 64 , and generalization to higher dimensions requires an improved network architecture.
We evaluate all methods on the problems of 4 × super-resolution and deblurring, on one validation set image from each of the 1000 ImageNet classes, following [38].
Table 1 shows that DDRM outperforms all baseline methods, in all metrics, and on both problems with only 20 steps.
The only exception to this is that SNIPS achieves better KID than DDRM in noiseless deblurring, but it requires 50 × more NFEs to do so.
Note that the runtime of all the tested methods is perfectly linear with NFEs, with negligible differences in time per iteration.
DGP and DDRM-CC use ground-truth class labels for the test images to aid in the restoration process, and thus have an unfair advantage.
DDRM’s appeal compared to previous methods becomes more substantial when significant noise is added to the measurements.
Under this setting, DGP, RED, and SNIPS all fail to produce viable results, as evident in Table 2 and Figure 4.
Since DDRM is fast, we also evaluate it on the entire
ImageNet validation set in Appendix F.
5.3 Qualitative Experiments
DDRM produces high quality reconstructions across all the tested datasets and problems, as can be seen in Figures 1 and 3, and in Appendix I. As it is a posterior sampling algorithm, DDRM can produce multiple outputs for the same input, as demonstrated in Figure 5.
Moreover, the unconditional
ImageNet diffusion models can be used to solve inverse problems on out-of-distribution images with general content.
In Figure 6, we show DDRM successfully restoring 256 × 256 images from
USC-SIPI [53] that do not necessarily belong to any ImageNet class (more results in Appendix I).
Conclusions
We have introduced DDRM, a general sampling-based linear inverse problem solver based on unconditional/class-conditional diffusion generative models as learned priors.
Motivated by variational inference, DDRM only requires a few number of NFEs ( e.g. , 20) compared to other samplingbased baselines ( e.g. , 1000 for SNIPS) and achieves scalability in multiple useful scenarios, including denoising, super-resolution, deblurring, inpainting, and colorization.
We demonstrate the empirical successes of DDRM on various problems and datasets, including general natural images outside the distribution of the observed training set.
To our best knowledge, DDRM is the first unsupervised method that effectively and efficiently samples from the posterior distribution of inverse problems with significant noise, and can work on natural images with general content.
In terms of future work, apart from further optimizing the timestep and variance schedules, it would be interesting to investigate the following: ( i ) applying DDRM to non-linear inverse problems, ( ii ) addressing scenarios where the degradation operator is unknown, and ( iii ) self-supervised training techniques inspired by DDRM as well as ones used in supervised techniques [41] that further improve performance of unsupervised models for image restoration.
Acknowledgements
We thank Kristy Choi, Charlie Marx, and Avital Shafran for insightful discussions and feedback.
This research was supported by NSF (#1651565, #1522054, #1733686), ONR (N00014-19-1-2145),
AFOSR (FA9550-19-1-0024), ARO (W911NF-21-1-0125), Sloan Fellowship, Amazon AWS, Stanford Institute for Human-Centered Artificial Intelligence (HAI), Google Cloud, the Israel Science
Foundation (ISF) under Grant 335/18, the Israeli Council For Higher Education - Planning & Budgeting Committee, and the Stephen A. Kreynes Fellowship.