Improving Diffusion Models for Inverse Problems using Manifold Constraints
Abstract
Recently, diffusion models have been used to solve various inverse problems in an unsupervised manner with appropriate modiﬁcations to the sampling process.
However, the current solvers, which recursively apply a reverse diffusion step followed by a projection-based measurement consistency step, often produce suboptimal results.
By studying the generative sampling path, here we show that current solvers throw the sample path off the data manifold, and hence the error accumulates.
To address this, we propose an additional correction term inspired by the manifold constraint, which can be used synergistically with the previous solvers to make the iterations close to the manifold.
The proposed manifold constraint is straightforward to implement within a few lines of code, yet boosts the performance by a surprisingly large margin.
With extensive experiments, we show that our method is superior to the previous methods both theoretically and empirically, producing promising results in many applications such as image inpainting, colorization, and sparse-view computed tomography.
Code available here
Introduction
Diffusion models have shown impressive performance both as generative models themselves [41, 13], and also as unsupervised inverse problem solvers [41, 8, 9, 25] that do not require problem-speciﬁc training.
Speciﬁcally, given a pre-trained unconditional score function (i e. denoiser), solving the reverse stochastic differential equation (SDE) numerically would amount to sampling from the data generating distribution [41].
For many different inverse problems (e g. super
resolution [8, 9], inpainting [41, 9], compressed-sensing MRI (CS-MRI) [40, 9], sparse view CT (SV-CT)
[40], etc.)
, it was shown that simple incorporation of the measurement process produces satisfactory conditional samples, even when the model was not trained for the speciﬁc problem.
Nevertheless, for certain problems (e.g. inpainting), currently used algorithms often produce unsatisfactory results when implemented naively (e.g. boundary artifacts, as shown in Fig.
1 (b)).
The authors in [32] showed that in order to produce high quality reconstructions, one needs to iterate back and forth between the noising and the denoising step at least > 10 times per iteration .
These iterations are computationally demanding and should be avoided, considering that diffusion models are slow to sample from even without such iterations.
On the other hand, a classic result of Tweedie’s formula [37, 42] shows that one can perform Bayes optimal denoising in one step, once we know the gradient of the log density.
Extending such result, it was recently shown that one can indeed perform a single-step denoising with learned score functions for denoising problems from the general exponential family [28].
In this work, we leverage the denoising result through Tweedie’s formula and show that such denoised samples can be the key to signiﬁcantly improving the performance of reconstruction using diffusion models across arbitrary linear inverse problems, despite the simplicity in the implementation.
Moreover, we theoretically prove that if the score function estimation is globally optimal, the correction term from the manifold constraint enforces the sample path to stay on the plane tangent to the data manifold 1 , so by combining with the reverse diffusion step, the solution becomes more stable and accurate.
Related Works 2.1 Diffusion Models
Continuous Form For a continuous diffusion process x ( t ) ∈ R n , t ∈ [0 , 1] , we set x (0) ∼ p 0 ( x ) =
p data , where p data represents the data distribution of interest, and x (1) ∼ p 1 ( x ) , with p 1 ( x ) approximating spherical Gaussian distribution, containing no information of data.
Here, the forward noising process is deﬁned with the following It ˆ o stochastic differential equation (SDE) [41]:
[Equation]
with ¯ f : R d (cid:55)→ R d deﬁning the linear drift function, ¯ g ( t ) : R (cid:55)→
R deﬁning a scalar diffusion coefﬁcient, and w ∈ R n denoting the standard n − dimensional Wiener process.
The forward SDE in (1) is coupled with the following reverse SDE by the Anderson’s theorem [1, 41]:
[Equation]
with dt denoting the inﬁnitesimal negative time step, and ¯ w deﬁning the standard Wiener process running backward in time.
Note that the reverse SDE deﬁnes the generative process through the score function ∇ x log p t ( x ) , which in practice, is typically replaced with ∇ x log p 0 t ( x ( t ) | x (0)) to minimize the following denoising score-matching objective
[Equation]
Once the parameter θ ∗ for the score function is estimated, one can replace the score function in (2) with s θ ∗ ( x ( t ) , t ) to solve the reverse SDE [41].
Discrete Form Due to the linearity of ¯ f and ¯ g , the forward diffusion step can be implemented with a simple reparameterization trick [29].
Namely, the general form of the forward diffusion is
[Equation]
where we have replaced the continuous index t ∈ [0 , 1] with the discrete index i ∈ N .
On the other hand, the discrete reverse diffusion step can in general be represented as
[Equation]
where we have replaced the ground truth score function with the trained one.
We detail the choice of a i , b i , f , g in Appendix.
B.
2.2 Conditional Generative models for Inverse problems
The main problem of our interest in this paper is the inverse problem, retrieving the unknown x ∈ R n from a measurement y :
[Equation]
where (cid:15) ∈
R m is the noise in the measurement.
Accordingly, for the case of the inverse problems, our goal is to generate samples from a conditional distribution with respect to the measurement y , i.e. p ( x | y ) .
Accordingly, the score function ∇ x log p t ( x ) in (2) should be replaced by the conditional score ∇ x log p t ( x | y ) .
Unfortunately, this strictly restricts the generalization capability of the neural network since the conditional score should be retrained whenever the conditions change.
To address this, recent conditional diffusion models [22, 41, 8, 9] utilize the unconditional score function ∇ x log p t ( x ) but rely on a projection-based measurement constraint to impose the conditions.
Speciﬁcally, one can apply the following:
[Equation]
where A , b i are functions of H , y , and x 0 .
Note that (7) is identical to the unconditional reverse diffusion step in (5), whereas (8) effectively imposes the condition.
It was shown in [9] that any general contraction mapping (e.g. projection onto convex sets, gradient step) may be utilized as (8) to impose the constraint.
Another recent work [25] advancing [26] establishes the state-of-the-art (SOTA) in solving noisy inverse problems with unconditional diffusion models, by running the conditional reverse diffusion process in the spectral domain achieved by performing singular value decomposition (SVD), and leveraging approximate gradient of the log likelihood term in the spectral space.
The authors show that feasible solutions can be obtained with as small as 20 diffusion steps.
Prior to the development of diffusion models, Plug-and-Play (PnP) models [47, 53, 44] were used in a similar fashion by utilizing a general-purpose unconditional denoiser in the place of proximal mappings in model-based iterative reconstruction methods [5, 3].
Similarly, outside the context of diffusion models, iterative denoising followed by projection-based data consistency was proposed in [44].
In such view, diffusion models can be understood as generative variant of PnPs trained with multiple scales of noise.
GAN-based solvers are also widly explored [4, 10, 20], where the pre-trained generators are tuned at the test time by optimizing over the latent, the parameters, or jointly.
2.3 Tweedie’s formula for denoising
In the case of Gaussian noise, a classic result of Tweedie’s formula [37] tells us that one can achieve the denoised result by computing the posterior expectation:
[Equation]
where the noise is modeled by ˜ x ∼ N ( x , σ 2 I ) .
If we consider a diffusion model in which the forward step is modeled as x i ∼ N ( a i x 0 , b
2 i I ) (discrete form), the Tweedie’s formula can be rewritten as:
[Equation]
Tweedie’s formula is in fact not only relevant to Gaussian denoising in the Bayesian framework, but have also been extended to be in close relation with kernel regression [34].
Moreover, it was shown that it can be applied to arbitrary exponential noise distributions beyond Gaussian [14, 28].
In the following, we use this key property to develop our algorithm.
Conditional Diffusion using Manifold Constraints
Although our original motivation of using the measurement constraint step in (8) was to utilize the unconditionally trained score function in the reverse diffusion step in (7), there is room for imposing additional constraints while still using the unconditionally trained score function.
Speciﬁcally, the Bayes rule p ( x | y )
=
p ( y | x ) p ( x ) /p
( y ) leads to
[Equation]
Hence, the score function in the reverse SDE in (7) can be replaced by (11), leading to
[Equation]
where α and W depend on the noise covariance, if the noise (cid:15) in (6) is Gaussian.
Now, one of the important contributions of this paper is to reveal that the Bayes optimal denoising step in (10) from the Tweedie’s formula leads to a preferred condition both empirically and theoretically.
Speciﬁcally, we deﬁne the set constraint for x i , called the manifold constrained gradient (MCG) , so that the gradient of the measurement term stays on the manifold (see Theorem 1):
[Equation]
∈
[Equation]
i , where X
[Equation]
To deal with the potential deviation from the measurement consistency, we again impose the data consistency step (8).
Putting them together, the discrete reverse diffusion under the additional manifold constraint and the data consistency can be represented by
[Equation]
We illustrate our scheme visually in Fig.
1 (a), speciﬁcally for the task of image inpainting.
The additional step leads to a dramatic performance boost, as can be seen in Fig.
1 (b).
Note that while the mapping (10) does not rely on the measurement, our gradient term in (14) incorporates the information of y so that the gradient of the measurement terms stays on the manifold.
In the following, we study the theoretical properties of the method.
Further algorithmic details and adaptations to each problem that we tackle are presented in Section C.
We note that the authors of [19] proposed a similar gradient method for the application of temporal imputation and super-resolution.
When combining (14) with (15), one can arrive at a similar gradient method proposed in [19], and hence our method can be seen as a generalization to arbitrary linear inverse problems.
Furthermore, there are vast literature in the context of PnP models that utilize pretrained denoisers together with gradient of the log-likelihood to solve inverse problems [30, 48, 11].
Among them, [30] is especially relevant to this work since their method relies on modiﬁed Langevin diffusion, together with Tweedie’s denoising and projections to the measurement subspace.
Geometry of Diffusion Models and Manifold Constrained Gradient
In this section, we theoretically support the effectiveness of the proposed algorithm by showing the problematic behavior of the earlier algorithm and how the proposed algorithm resolves the problem.
We defer all proofs in the supplementary section.
To begin with, we borrow a geometrical viewpoint of the data manifold.
Notation For a scalar a , points x , y and a set A , we use the following notations.
aA :=
{ a x : x ∈
A } ; d ( x , A ) :=
inf y ∈
A || x − y || 2 ; B r ( A ) :
= { x : d ( x , A ) < r } ; T x M : the tangent space to a manifold M at x ; J f : the Jacobian matrix of a vector valued function f .
We deﬁne p
= p data .
To develop the theory, we need an assumption on the data distribution.
Assumption 1 (Strong manifold assumption: linear structure) .
Suppose M ⊂ R n is the set of all data points, here we call the data manifold.
Then, the manifold coincides with the tangent space with dimension l (cid:28) n .
[Equation]
Moreover, the data distribution p 0 is the uniform distribution on the data manifold M .
[Figure]
We need to recall that the conventional manifold assumption is about the intrinsic geometry of data points having a low dimensional nature.
However, we assume more in this work: the manifold is locally linear.
Although this stronger assumption might narrow the practice of the theory, the geometric approach may provide new insights on diffusion models.
Under this assumption, the following proposition shows how the data perturbed by noise lies in the ambient space, illustrated pictorially in Fig.
2a.
Proposition 1 (Concentration of noisy data) .
Consider the distribution of noisy data p i ( x i )
= (cid:82) p ( x i | x ) p 0 ( x ) d x , p ( x i | x ) ∼ N ( a i x , b 2 i I ) .
Then p i ( x i ) is concentrated on ( n − 1)
-dim manifold
M i := { y ∈ R n : d ( y , a i M ) =
r i :=
b i √ n − l } .
Rigorously, p i ( B (cid:15)r i ( M i )) > 1 − δ , for some small (cid:15), δ > 0 .
Remark 1 (Geometric interpretation of the diffusion process) .
Considering Proposition 1, the manifolds of noisy data can be interpreted as interpolating manifolds between the two: the hypersphere, where pure noise N ( a ∞ x 0 , b 2 ∞ ) is concentrated, and the clean data manifold.
In this regard, the diffusion steps are mere transitions from one manifold to another and the diffusion process is a transport from the data manifold to the hypersphere through interpolating manifolds.
See Fig.
2a.
Remark 2.
We can infer from the proposition that the score functions are trained only with the data points concentrated on the noisy data manifolds.
Therefore, inaccurate inference might be caused by application of a score function on points away from the noisy data manifold.
Proposition 2 (score function) .
Suppose s θ is the minimizer of the denoising score matching loss in (3) .
Let Q i be the function that maps x
i to ˆ x 0 for each i ,
[Equation]
Then, Q i ( x i ) ∈
M and J 2 Q i =
J Q i =
J TQ i : R d → T Q i ( x i ) M .
Intuitively, Q i is locally an orthogonal projection onto M .
According to the proposition, the score function only concerns the normal direction of the data manifold.
In other words, the score function cannot discriminate two data points whose difference is tangent to the manifold.
In solving inverse problems, however, we desire to discriminate data points to reconstruct the original signal, and the discrimination is achievable by measurement ﬁdelity.
In order to achieve the original signal, the measurement plays a role in correcting the tangent component near the data manifold.
Furthermore, with regard to remark 2, diffusion model-based inverse problem solvers should follow the tangent component.
The following theorem shows how existing algorithms and the proposed method are different in this regard.
Theorem 1 (Manifold constrained gradient) .
A correction by the manifold constrained gradient does not leave the data manifold.
Formally,
[Equation]
the gradient is the projection of the data ﬁdelity term onto T ˆ x 0 M ,
This theorem suggests that in diffusion models, the naive measurement ﬁdelity step (without considering the data manifold) pushes the inference path out of the manifolds and might lead to inaccurate reconstruction.
(To see this pictorially, see section.
D, and Fig.
On the other hand, our correction term from the manifold constraint guides the diffusion to lie on the data manifold, leading to better reconstruction.
Such geometric views are illustrated in Fig. 2b.
Remark 3.
One may concern that the suboptimality of the denoising score matching loss optimization may lead to inaccurate inference of the MCG steps.
In practice, however, most of the error in denoising score matching is concentrated on t ∼ 1 [9], and in such region, the Tweedie’s inference cannot make meaningful images.
That is, the score function cannot detect the data manifold.
Nonetheless, in this regime, the magnitudes of the MCGs are small when the denoising score is inaccurate, and hence the matters arising from suboptimality is minimal.
As t → 0 , the estimation becomes exact, and subsequently leads to accurate implementation of the MCG.
Experiments
For all tasks, we aim to verify the superiority of our method against other diffusion model-based approaches, and also against strong supervised learning-based baselines.
Further details can be found in Section.
F.
Datasets and Implementation For inpainting, we use FFHQ 256 × 256 [24], and ImageNet 256 × 256 [12] to validate our method.
We utilize pre-trained models from the open sourced repository based on the implementation of ADM (VP-SDE)
We validate the performance onheld-out validation set images for both FFHQ and ImageNet dataset.
For the colorization task, we use FFHQ 256 × 256, and LSUN-bedroom 256 × 256 [51].
We use pre-trained score functions from
[Table] [Table]
Inpainting Score-SDE [41], REPAINT [32],
DDRM [25] were chosen as baseline diffusion models to compare against the proposed method.
For a fair comparison, we use the same score function for all methods including MCG, and only differentiate the inference method that is used.
Another class of generative models: GAN-based inverse problem solver, IAGAN [20] is considered as a comparison method for FFHQ speciﬁcally.
We also include comparisons against supervised learning based baselines: LaMa [43], AOT-GAN [52], ICT [49], and DSI [35].
We use various forms of inpainting masks: box (128 × 128 sized square region is missing 2 ), extreme (only the box region is existent), random (90-95% of pixels are missing), and LaMa-wide.
Quantitative evaluation is performed with two metrics - Frechet Inception Distance (FID)-1k [17], and Learned Perceptual Image Patch Similarity (LPIPS) [54].
[Table]
Our method outperforms the diffusion model baselines [41, 32, 25] by a large margin.
Moreover, our method is also competitive with, or even better than the best-in-class fully supervised methods, as can be seen in Table 1.
In
Fig.
3, we depict representative results that show the superiority of the method, where we see in both the box-type and random dropping that
MCG performs very well on all experiments.
Colorization We choose score-SDE [41], and
DDRM [25] as diffusion-model based comparison methods, and also compare against cINN [2], and pix2pix [21].
Two metrics were used for evaluation: structural similarity index (SSIM), and LPIPS.
Consistent with the ﬁndings from inpainting, we achieve much improved performance than score-SDE, and also is favorable against state-of-the-art (SOTA) superivsed learning based methods.
In Table 2, we see that the proposed method outperforms all other methods in terms of both PSNR/LPIPS in LSUN-bedroom, and also achieves strong performance in the colorization of FFHQ dataset.
CT reconstruction To the best of our knowledge, [40] is the only method that tackles CT reconstruction directly with diffusion models.
We compare our method against [40], which we refer to as score-CT henceforth.
We also compare with the best-in-class supervised learning methods, cGAN [15] and SIN-4c-PRN [50].
As a compressed sensing baseline, FISTA-TV [3] was included, along with the analytical reconstruction method, FBP.
We use two standard metrics - peak-signalto-noise-ratio (PSNR), and SSIM for quantitative evaluation.
From Table 3, we see that the newly proposed MCG method outperforms the previous score-CT [40] by a large margin.
We can observe the superiority of MCG over other methods more clearly in Fig.
4, where MCG reconstructs the measurement with high ﬁdelity and detail.
All other methods including the fully supervised baselines fall behind the proposed method.
[Table]
Ablation studies We perform three ablation studies: 1)
As both the MCG term and the projection term contain information about the measurement y , we observe the contribution of each term to the ﬁxed solution.
To further clarify the efﬁcacy of the gradient step combined with Tweedie’s denoising, we also consider the case where the gradient of the log likelihood is computed not in the noiseless regime, but in the noise level matching the current iteration.
Speciﬁcally, we deﬁne x (cid:48)
i − 1 :
= f
( x i , s θ )
+ g ( x i ) z , z ∼
N (0 , I ) , y i − 1 ∼ p
( y i − 1 | y 0 ) , and implement the gradient step as ∇ x i (cid:107) y
i − 1 − Hx (cid:48)
i − 1 (cid:107) 22 . 2)
As the performance of diffusion models depend heavily on the number of NFEs, we observe the trade-off of each diffusion model when varying the NFE from 20 to 1000.
Moreover, for completeness, we measure the runtime of each algorithms including the non-diffusion based methods in wall-clock time computed with a commodity GPU in
Table.
4. 3) Setting α = 0 .
0 reduces our method to [9].
We show the difference in the performance by varying the values of α .
[Table]
First, we see in Table.
5 that using only the MCG step leads to improved performance in terms of
LPIPS, but introduces error in the measurement consistency (measured with MSE).
Combining both the projection and MCG leads to perfect data consistency along with further improved
[Figure]
reconstruction.
When considering gradient steps without Tweedie’s denoising (i.e. keeping the noise level at the i th step), the performance heavily degrades, especially when implemented without the projection steps.
Here, we see that the proposed denoising step to utilize ˆ x 0 is indeed the key to the superior performance.
Second, looking at Fig.
5a, we immediately see that the graph of MCG stays in the lowest (best)
LPIPS regime across all NFEs by a large margin, except for when the NFE drops below 100.
Here,
DDRM [25] takes over the 1st place - allegedly due to the DDIM sampling strategy they take.
The performance of RePAINT deteriorates rapidly as we decrease NFE.
Furthermore, we observe that the LPIPS of score-SDE [41] actually increases (i.e. worsen), as we increase the number of NFEs from a few hundred to one thousand.
This suggests that the inference process that score-SDE takes (i.e. projection only) is inherently ﬂawed, and cannot be corrected by taking small enough steps.
In
Table.
4, we list the runtime of all the methods that were used for comparison in the task of inpainting.
Note that the proposed method takes longer for compute than score-SDE albeit having the same
NFE.
The gap is due to the backpropagation steps that are required for the MCG step, where the gap can be potentially ameliorated by switching to JAX [6] implementation from the current PyTorch implementation.
Lastly, we observe the difference in the performance as we vary the values of α .
Implementation-wise,
we ﬁnd that we yield superior results when normalizing the squared norm with the norm of itself (e.g. α = α (cid:48) / (cid:107) W ( y − H ˆ x 0 ) (cid:107) , where α (cid:48) is some constant).
In order to avoid cluttered notation, we instead experiment with changing the values of α (cid:48) in Fig.
5b.
Inspecting Fig.
5b, we see that α values within the range [0 . 1 , 1 . 0] produce satisfactory results.
α values that are too low do not fully enjoy the advantages of MCG and collapses to the projection-only method, while using too high values of α results in exploding gradients, and the reconstruction saturates.
Properties of our method Our proposed method is fully unsupervised and is not trained on solving a speciﬁc inverse problem.
For example, our box masks and random masks have very different forms of erasing the pixel values.
Nevertheless, our method generalizes perfectly well to such different measurement conditions, while other methods have a large performance gap between the different mask shapes.
We further note two appealing properties of our method as an inverse problem solver: 1) the ability to generate multiple solutions given a condition, and 2) the ability to maintain perfect measurement consistency.
The former ability often lacks in supervised learningbased methods [43, 50], and the latter is often not satisﬁed for some unsupervised GAN-based solutions [10, 4].
Conclusion
In this work, we proposed a general framework that can greatly enhance the performance of the diffusion model-based solvers for solving inverse problems.
We showed several promising applications - inpainting, colorization, sparse view CT reconstruction, and showed that our method can outperform the current state-of-the-art methods.
We analyzed our method theoretically and show that
MCG prevents the data generation process from falling off the manifold, thereby reducing the errors that might accumulate at every step.
Further, we showed that MCG controls the direction tangent to the data manifold, whereas the score function controls the direction that is normal, such that the two components complement each other.
Limitations and Broader Impact
The proposed method is inherently stochastic since the diffusion model is the main workhorse of the algorithm.
When the dimension m is pushed to low values, at times, our method fails to produce high quality reconstructions, albeit being better than the other methods overall.
For extreme cases of inpainting (e.g. Half masks) with the ImageNet model, we often observe artifacts in our reconstruction (e.g. generating perfectly symmetric images), which we discuss in further detail in Sec.
E. We note that our method is slow to sample from, inheriting the existing limitations of diffusion models.
This would likely beneﬁt from leveraging recent solvers aimed at accelerating the inference speed of diffusion models.
In line with the arguments of other generative model-based inverse problem solvers, our method is a solver that relies heavily on the underlying diffusion model, and can thus potentially create malicious content such as deepfakes.
Further, the reconstructions could intensify the social bias that is already existent in the training dataset.
Acknowledgments and Disclosure of Funding
This research was supported by Field-oriented Technology Development Project for Customs Administration through National Research Foundation of Korea(NRF) funded by the Ministry of Science &
ICT and Korea Customs Service (NRF-2021M3I1A1097938, NRF-2021M3I1A1097910), by the
Korea Health Technology R&D Project through the Korea Health Industry Development Institute (KHIDI), which is funded by the Ministry of Health & Welfare, Republic of Korea (grant number:
[-where we applied v Tn u t = 
-]
HU21C0222), and by the KAIST Key Research Institute (Interdisciplinary Research Group) Project.
[-The sampler is based on basic ancestral sampling (AS) of [18], and the default conﬁguration requires N = 1000 , α = 1 . 
-]