Improving Diffusion Models for Inverse Problems using Manifold Constraints
Anonymous Author(s)
AfﬁliationAddress
Abstract
Recently, diffusion models have been used to solve various inverse problems inan unsupervised manner with appropriate modiﬁcations to the sampling process.
However, the current solvers, which recursively apply a reverse diffusion stepfollowed by a projection-based measurement consistency step, often produce suboptimal results.
By studying the generative sampling path, here we show thatcurrent solvers throw the sample path off the data manifold, and hence the erroraccumulates.
To address this, we propose an additional correction term inspiredby the manifold constraint, which can be used synergistically with the previoussolvers to make the iterations close to the manifold.
The proposed manifoldconstraint is straightforward to implement within a few lines of code, yet booststhe performance by a surprisingly large margin.
With extensive experiments,we show that our method is superior to the previous methods both theoreticallyand empirically, producing promising results in many applications such as imageinpainting, colorization, and sparse-view computed tomography.
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ﬁctraining.
Speciﬁcally, given a pre-trained unconditional score function (i e. denoiser), solving thereverse stochastic differential equation (SDE) numerically would amount to sampling from the datagenerating 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 conditionalsamples, 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)).
Theauthors in [32] showed that in order to produce high quality reconstructions, one needs to iterateback and forth between the noising and the denoising step at least > 10 times per iteration .
Theseiterations are computationally demanding and should be avoided, considering that diffusion modelsare slow to sample from even without such iterations.
On the other hand, a classic result of Tweedie’sformula [37, 42] shows that one can perform Bayes optimal denoising in one step, once we know thegradient of the log density.
Extending such result, it was recently shown that one can indeed performa single-step denoising with learned score functions for denoising problems from general exponentialfamily noises [28].
In this work, we leverage the denoising result through Tweedie’s formula and show that such denoisedsamples can be the key to signiﬁcantly improve the performance of reconstruction using diffusionmodels across arbitrary inverse problems, despite the simplicity in the implementation.
Moreover, we theoretically prove that if the score function estimation is globally optimal, the correction term fromthe 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 ) , withp 1 ( x ) approximating spherical Gaussian distribution, containing no information of data.
Here, theforward noising process is deﬁned with the following It ˆ o stochastic differential equation (SDE) [41]:
[Equation]
(1) with ¯ f : R d (cid:55)→ R d deﬁning the linear drift function, ¯ g ( t ) : R (cid:55)→
R deﬁning a scalar diffusioncoefﬁ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 Anderson’s theorem [1, 41]
[Equation]
with dt denoting the inﬁnitesimal negative time step, and ¯ w deﬁning the standard Wiener processrunning backward in time.
Note that the reverse SDE deﬁnes the generative process through thescore function ∇ x log p t ( x ) , which in practice, is typically replaced with ∇ x log p 0 t ( x ( t ) | x (0)) tominimize 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.
Discrete Form Due to the linearity of ¯ f and ¯ g , the forward diffusion step can be implemented witha 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 otherhand, the discrete reverse diffusion step is implemented as
[Equation]
where we have replaced the ground truth score function with the trained one.
We detail the choice ofa i , b i , f , g in Section.
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 nfrom 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 measurementy , i.e. p ( x | y ) .
Accordingly, the score function ∇ x log p t ( x ) in (2) should be replaced by theconditional score ∇ x log p t ( x | y ) .
Unfortunately, this strictly restricts the generalization capabilityof 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 scorefunction ∇ 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 function of H , y 0 , and x 0 .
Note that (7) is identical to the unconditional reversediffusion step in (5), whereas (8) effectively imposes the condition.
It was shown in [9] that anygeneral contraction mapping (e.g. projection onto convex sets, gradient step) may be utilized as (8) toimpose the constraint.
Another recent work [25] advancing [26] establishes the state-of-the-art (SOTA) in solving noisyinverse problems with unconditional diffusion models, by running the conditional reverse diffusionprocess in the spectral domain achieved by performing singular value decomposition (SVD), andleveraging approximate gradient of the log likelihood term in the spectral space.
The authors showthat feasible solutions can be obtained with as small as 20 diffusion steps.
Prior to the developmentof diffusion models, Plug-and-Play (PnP) models [47, 53, 44] were used in a similar fashion byutilizing a general-purpose unconditional denoiser in the place of proximal mappings in model-basedreconstruction methods [5, 3].
Similarly, outside the context of diffusion models, iterative denoisingfollowed by projection-based data consistency was proposed in [44].
In such view, diffusion modelscan be understood as generative variant of PnPs trained with multiple scales of noise.
GAN-basedsolvers are also popular [4, 10, 20], where the pre-trained generators are tuned at test time byoptimizing 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 achievethe denoised result by computing posterior expectation:
[Equation]
when the noise is modeled by ˜ x ∼ N ( x , σ 2 I ) .
If we consider a diffusion model of which forwardstep is modeled as x i ∼ N ( a i x 0 , b
2 i I ) (discrete form), the Tweedie’s formular is rewritten:
[Equation]
Tweedie’s formula is in fact not only relevant to Gaussian denoising in the Bayesian framework, buthave also been extended to be in close relation with kernel regression [34].
Moreover, it was shownthat it can be applied to arbitrary exponential noise distributions beyond Gaussian [14, 28].
In thefollowing, 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 theunconditionally trained score function in the reverse diffusion step in (7), there is room for imposingadditional 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 (2) 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 denoisingstep in (10) from Tweedie formula leads to a preferred condition both empirically and theoretically.
Speciﬁcally, we deﬁne the set constraint for x i , called the manifold constraint gradient (MCG) , sothat the gradient of the measurement terms stays on the manifold (see Theorem 1):
[Equation]
, where X
[Equation]
To deal with the potential deviation from the measurement consistency, we again impose the dataconsistency step (8).
Putting them together, the discrete reverse diffusion under the additionalmanifold 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.
Theadditional step leads to dramatic performance boost, as can be seen in Fig.
1 (b).
Note that while themapping (10) does not rely on the measurement, our gradient term in (14) incorporates the informationof y so that the gradient of the measurement terms stays on the manifold.
In the following, wearticulate how the proposed MCG can be incorporated into each application, and study the theoreticalproperties of the method.
Further algorithmic details and adaptations to the each problem that wetackle are presented in Section C.
We note that the authors of [19] proposed a similar gradient method for the application of temporalimputation and super-resolution.
When combining (14) with (15), one can arrive at a similar gradientmethod proposed in [19], and hence our method can be seen as a generalization to arbitrary linearinverse 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 Langevindiffusion, together with Tweedie’s denoising and projections to the measurement subspace.
Geometry of Diffusion Model and Manifold Constraint Gradient
In this section, we theoretically support the effectiveness of the proposed algorithm by showing theproblematic 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 viewpointof 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 of to amanifold 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 .
We need to recall that the conventional manifold assumption is about the intrinsic geometry of datapoints having a low dimensional nature.
However, we assume more in this work: the manifoldis locally linear.
Although this stronger assumption might narrow the practice of the theory, thegeometric approach may provide new insights on diffusion models.
Under this assumption, thefollowing proposition shows how the data perturbed by noise lies in the ambient space, illustratedpictorially in Fig.
2a.
[Figure]
Proposition 1 (Concentration of noisy data) .
Consider the distribution of noisy data p i ( x i )
= 141 (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 ( B (cid:15)r i ( M i )) > 1 − δ .
Remark 1 (Geometric interpretation of diffusion process) .
Considering Proposition 1, the manifoldsof noise data can be interpreted as interpolating manifolds between the two: the hypersphere, wherepure noise N ( x 0 , b 2 ∞ ) is concentrated, and the clean data manifold.
In this regard, the diffusionsteps are mere transitions from one manifold to another and the diffusion process is a transport fromthe 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 datapoints concentrated on the noisy data manifolds.
Therefore, inaccurate inference might be caused byapplication 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 152 orthogonal projection onto M .
According to the proposition, the score function only concerns the normal direction of the datamanifold.
In other words, the score function cannot discriminate two data points whose difference istangent to the manifold.
In solving inverse problems, however, we desire to discriminate data pointsto reconstruct the original signal, and the discrimination is achievable by measurement ﬁdelity.
Inorder to achieve the original signal, the measurement plays a role in correcting the tangent componentnear the data manifold.
Furthermore, with regard to remark 2, diffusion model-based inverse problemsolvers should follow the tangent component.
The following theorem shows how existing algorithmsand the proposed method are different in this regard.
Theorem 1 (Manifold constrained gradient) .
A correction by the manifold constrained gradient doesnot leave the data manifold.
Formally,
[Equation]
the gradient is the projection of 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 inaccuratereconstruction.
(To see this pictorially, see section.
D, and Fig.
On the other hand, our correctionterm from the manifold constraint guides the diffusion to lie on the data manifold, leading to betterreconstruction.
Such geometric views are illustrated in Fig. 2b.
Remark 3.
One may concern that the suboptimality of the denoising score matching loss optimizationmay lead to inaccurate inference of the MCG steps.
In practice, however, most of the error in denoisingscore matching is concentrated on t ∼ 1 [9], and in such region, the Tweedie’s inference cannot makemeaningful images.
That is, the score function cannot detect the data manifold.
Nonetheless, in thisregime, the magnitudes of the MCGs are small when the denoising score is inaccurate, and hencethe matters arising from suboptimality is minimal.
As t → 0 , the estimation becomes exact, andsubsequently leads to accurate implementation of the MCG.
Experiments
For all tasks, we aim to verify the superiority of our method against other diffusion model-basedapproaches, and also against strong supervised learning-based baselines.
Further details can be foundin Section.
F.
Datasets and Implementation For inpainting, we use FFHQ 256 × 256 [24], and ImageNet256 × 256 [12] to validate our method.
We utilize pre-trained models from the open sourced repositorybased on the implementation of ADM (VP-SDE)
We validate the performance on 1000 held-outvalidation set images for both FFHQ and ImageNet dataset.
For the colorization task, we use FFHQ,and LSUN-bedroom 256 × 256 [51].
We use pre-trained score functions from score-SDE [41] basedon VE-SDE.
We use 300 validation images for testing the performance with respect to the LSUNbedroom dataset.
For experiments with CT, we train our model based on ncsnpp as a VE-SDEfrom score-SDE [41], on the 2016 American Association of Physicists in Medicine (AAPM) grandchallenge dataset, and we process the data as in [23].
Speciﬁcally, the dataset contains 3839 trainingimages resized to 256 × 256 resolution.
We simulate the CT measurement process with parallel beamgeometry with evenly-spaced 180 degrees.
Evaluation is performed on 421 held-out validation imagesfrom the AAPM challenge.
[Table] Inpainting Score-SDE [41], REPAINT [32], DDRM [25] were chosen as baseline diffusion modelsto compare against the proposed method.
For a fair comparison, we use the same score functionfor all methods including MCG, and only differentiate the inference method that is used.
Anotherclass of generative models: GAN-based inverse problem solver, IAGAN [20] is considered as acomparison method for FFHQ speciﬁcally.
We also include comparisons against supervised learningbased baselines: LaMa [43], AOT-GAN [52], ICT [49], and DSI [35].
We use various forms ofinpainting masks: box (128 × 128 sized square region is missing 2 ), extreme (only the box regionis existent), random (90-95% of pixels are missing), and LaMa-wide.
Quantitative evaluation isperformed with two metrics - Frechet Inception Distance (FID)-1k [17], and Learned Perceptual
Image Patch Similarity (LPIPS) [54].
Data [Table] Our method outperforms the diffusion modelbaselines [41, 32, 25] by a large margin.
Moreover, our method is also competitive with, oreven better than the best-in-class fully supervised methods, as can be seen in Table 1.
In Fig.
3, we depict representative results that showthe superiority of the method, where we seein 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], andpix2pix [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 scoreSDE, 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/LPIPSin LSUN-bedroom, and also achieves strong performance in the colorization of FFHQ dataset.
[Table]
CT reconstruction To the best of our knowledge, [40] is the only method that tackles CTreconstruction directly with diffusion models.
We compare our method against [40], which werefer to as score-CT henceforth.
We also compare with the best-in-class supervised learningmethods, cGAN [15] and SIN-4c-PRN [50].
Asa compressed sensing baseline, FISTA-TV [3]was included, along with the analytical reconstruction method, FBP.
We use two standardmetrics - peak-signal-to-noise-ratio (PSNR), and
SSIM for quantitative evaluation.
From Table 3,we see that the newly proposed MCG methodoutperforms the previous score-CT [40] by a large margin.
We can observe the superiority of MCGover 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 proposedmethod.
[Table] Ablation studies We perform three ablation studies: 1) As both the MCG term and the projection term containinformation about the measurement y , we observe thecontribution of each term to the ﬁxed solution.
To furtherclarify the efﬁcacy of the gradient step combined with
Tweedie’s denoising, we also consider the case where thegradient step is implemented without denoising (i.e. gradient w.r.t.
(cid:107) y i − Hx i (cid:107) 22 ). As the performance of diffusion models depend heavily on the number of NFEs, weobserve 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 thenon-diffusion based methods in wall-clock time computedwith a commodity GPU. 3) Setting α = 0 .
0 reduces ourmethod to [9].
We show the difference in the performanceby varying the values of α .
First, we see in Fig.
5a 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 boththe projection and MCG leads to perfect data consistency along with further improved reconstruction.
When considering gradient steps without Tweedie’s denoising (i.e. keeping the noise level at thei th step), the performance heavily degrades, especially when implemented without the projectionsteps.
Here, we see that the proposed denoising step to utilize ˆ x 0 is indeed the key to the superiorperformance.
Second, looking at Fig.
5b, 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.
Theperformance of RePAINT deteriorates rapidly as we decrease NFE.
Furthermore, observe that the
LPIPS of score-SDE [41] actually increases (i.e. worsen), as we increase the number of NFEs froma 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
[Figure] NFE. The gap is due to the backpropagation steps that are required for the MCG step, where the gapcan be potentially ameliorated by switching to JAX [6] implementation from the current PyTorchimplementation. Lastly, seeing Fig.
5c, 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-onlymethod, while using too high values of α results in exploding gradients, and the reconstructionsaturates.
Properties of Our Method Our proposed method is fully unsupervised and is not trained onsolving a speciﬁc inverse problem.
For example, our box masks and random masks have verydifferent forms of erasing the pixel values.
Nevertheless, our method generalizes perfectly well tosuch different measurement conditions, while other methods have a large performance gap betweenthe different mask shapes.
We further note two appealing properties of our method as an inverseproblem solver: 1) the ability to generate multiple solutions given a condition, and 2) the ability tomaintain 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-basedsolutions [10, 4].
Conclusion
In this work, we proposed a general framework that can greatly enhance the performance of thediffusion model-based solvers for solving inverse problems.
We showed several promising applications - inpainting, colorization, sparse view CT reconstruction, and showed that our method canoutperform 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 errorsthat might accumulate at every step.
Further, we showed that MCG controls the direction tangent tothe data manifold, whereas the score function controls the direction that is normal, such that the twocomponents complement each other.
Limitations and Broader Impact
The proposed method is inherently stochastic since the diffusionmodel is the main workhorse of the algorithm.
When the dimension m is pushed to low values, attimes, our method fails to produce high quality reconstructions, albeit being better than the othermethods overall.
For extreme cases of inpainting (e.g. Half masks) with the ImageNet model, weoften observe artifacts in our reconstruction (e.g. generating perfectly symmetric images), which wediscuss in further detail in Sec.
E.We note that our method is slow to sample from, inheriting theexisting limitations of diffusion models.
This would likely beneﬁt from leveraging recent solversaimed at accelerating the inference speed of diffusion models.
In line with the arguments of othergenerative model-based inverse problem solvers, our method is a solver that relies heavily on theunderlying 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 trainingdataset.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] We discuss the limitations in (6).
(c) Did you discuss any potential negative societal impacts of your work?
[Yes] We discuss
potential negative impacts in (6).(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes]
(b) Did you include complete proofs of all theoretical results?
[Yes] We provide all proofs of results in supplementary material.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes] We include all code for our experiments in the supplementary material.
We will release the code once the paper is published.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes]
(c) Did you report error bars (e g., with respect to the random seed after running experiments multiple times)?
[No] Due to our limited resources we do not have time to run multiple sets of experiments.
(d) Did you include the total amount of compute and the type of resources used (e g., type
of GPUs, internal cluster, or cloud provider)?
[Yes]
If you are using existing assets (e
g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes] We have cited theoriginal works that released the datasets.
(b) Did you mention the license of the assets?
[No] Licenses are standard and can be found
online.
(c) Did you include any new assets either in the supplemental material or as a URL?
[Yes]We include our implementation as the supplementary material.
We will release thecode upon publication.
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[No] All datasets used in our work are publicly available.
(e) Did you discuss whether the data you are using/curating contains personally identiﬁable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
Did you include the full text of instructions given to participants and screenshots, if
applicable
[Table]
any participant , with links to Institutional Review you
the estimated paid to participants and the total amount
[Table]
spent on participant compensation
A