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:
HU21C0222), and by the KAIST Key Research Institute (Interdisciplinary Research Group) Project.
A Proofs
First, we remind our notation and the assumption.
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 .
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 state our proofs below.
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 .
Proof.
Suppose that the data manifold is an l -dimensional linear subspace.
By rotation and translation, we safely assume that M = { x ∈ R n : x l +1 = x l +2 = · · · = x n = 0 } .
Then, we can simply write d ( x , M ) =
(cid:113) x
2 l +1 + · · · + x 2 n , and M i = { x ∈ R n : x 2 l +1 + · · · + x 2 n =
r 2 i } .
For a given point x (cid:48)
=
( x (cid:48) 1 , x (cid:48) 2 , . . . )
∈ M , we consider p ( x | x (cid:48) )
∼ N ( a i x (cid:48) , b 2 i I ) and obtain a concentration inequality independent to the choice of x (cid:48) .
We need the standard Laurent-Massart bound for a chi-square variable [31].
When X is a chi-square distribution with k degrees of freedom,
[Equation]
As x 2 l +1 b 2 i + · · · + x 2 n b 2 i is a chi-square distribution with n − l degrees of freedom, by substituting t = ( n − l ) ε (cid:48) in the above bound,
[Equation]
Note that the above inequality does not depend on x 1 , . . .
x l , thus the choice of x (cid:48) ∈ M .
As a result, by setting ε = min { 1 − (cid:112) 1 − 2 √ ε (cid:48) , (cid:112) 1 + 2 √ ε (cid:48) + 2 ε (cid:48) − 1 } and δ = 2 e − ( n − l ) ε (cid:48) ,
[Equation]
thus
[Equation]
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 .
Proof.
To minimize (3), or equivalently,
[Equation]
By differentiating the objective with respect to s θ ( x t , t ) , we have
[Equation]
where we used p ( x t | x ) p ( x )
=
p ( x , x t ) = p ( x t ) p ( x | x t ) , p ( x t ) > 0 , and (cid:82) p ( x | x t )
d x = 1 in each line.
Here, Q i ( x i ) =
(cid:82) x p ( x | x i )
d x is the weighted average vector of points on the data manifold as p ( x | x i ) is supported on the data manifold.
Combining it with the assumption that the manifold is linear, Q i ( x i ) ∈ M .
Considering the symmetry of p ( x | x i ) about x i , p ( x | x i ) is a radial function on M , centering around the nearest point to x i on M .
Hence, Q i ( x i ) shall be the nearest point to x i of all points on M .
Therefore, J Q i is the orthogonal projection onto T Q i ( x i ) M .
Stating more rigorously, let u =
u t + u n ∈ R n for u t ∈ T Q i ( x i ) M , u n ⊥ T Q i ( x i ) M .
Then, for a scalar s , Q i ( x i + s u ) =
Q i ( x i )+ s u t , as only tangent component to the manifold change the nearest point.
By differentiating with respect to s , we obtain J Q i u
= u t , thus J 2 Q i =
J Q i .
For another vector v = v t + v n with v t ∈ T Q
i
( x i ) M , v n ⊥ T Q i ( x i ) M ,
[Equation]
where we applied v Tn u t =
= u
Tn v t .
Therefore, J Q i is symmetric, i.e. J TQ i =
J Q i , which concludes this proof.
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 ,
Proof.
[Equation]
where d = − 2 H T W T W ( y − H ˆ x 0 ) .
The ﬁrst and second equality is given by the chain rule and the last line is by Proposition 2.
In Fig 6, we illustrate how the proposed algorithm beneﬁts from mixing the MCG step with the conventional POCS step.
Pushing the points to the tangent directions, we expect less deviation from the manifold which is attributed to POCS.
B Discrete forms of SDE
Here, we review the different types of SDEs and sampling algorithms that we use throughout the paper for completeness.
We assume that the time horizon [0 , 1] is linearly split up into N discretization segments, such that all intervals have the length 1 /N , if not speciﬁed otherwise.
B.1 Forward diffusion
Due to the linearity of the drift and diffusion functions, we can analytically sample from p ( x i | x 0 ) via reparameterization trick:
[Equation]
In VP-SDE [18], one deﬁnes a linearly increasing noise schedule β 1 , β 2 , . .
. , β N ∈ (0 , 1) .
Further, we deﬁne α i = 1 − β i , and ¯ α i = (cid:81) ij =1 α j .
Then, the forward diffusion process can be implemented as
[Equation]
In VE-SDE [41], one deﬁnes a geometrically increasing noise schedule σ i = σ 0 (cid:16)
[Equation]
i − 1 N − 1 .
Since the drift function is zero, the forward diffusion simply becomes Brownian motion.
Concretely,
[Equation]
B.2 Reverse diffusion
First, for the case of VP-SDE, the reverse diffusion step is implemented by
[Equation]
where z θ ( x i , i ) is trained with the epsilon-matching scheme as in [18], and ˜ σ i is set to a learnable parameter as in [13].
Note that eq.
(19) was written in terms of z θ ( x i , i ) and not in terms of the score function, s θ ( x i , i ) .
One can re-write the expression using the relation z θ ( x i , i )
= −√
[Equation]
− ¯ α
i s θ ( x i , i ) , as
[Equation] [Equation]
Next, for the VE-SDE, the reverse diffusion step using the Euler-Maruyama solver [38] is given as
[Equation]
Summary is presented in Table 6.
[Table] [Table]
P
[Table]
{ α i } Ni =
[Table]
{ ˜ σ i } Ni =
[Table]
s θ
[Table]
α
[Table]
Data consistency 14: end for 15: return x
C Algorithms
Inpainting The forward model for inpainting is given as
[Equation]
where P ∈ { 0 , 1 } m × n is the matrix consisting of the columns with standard coordinate vectors indicating the indices of measurement.
For the steps in (14), (15), we choose the following
[Equation]
Speciﬁcally, A takes the orthogonal complement of x (cid:48)
i − 1 , meaning that the measurement subspace is corrected by y i , while the orthogonal components are updated from x (cid:48)
i − 1 .
Note that we use y i sampled from y to match the noise level of the current estimate.
We provide the algorithm used for inpainting in Algorithm.
The sampler is based on basic ancestral sampling (AS) of [18], and the default conﬁguration requires N = 1000 , α = 1 .
/ (cid:107) y − P
ˆ x 0 (cid:107) for sampling.
Colorization The forward model for colorization is speciﬁed as
[Equation]
where P is the matrix that was used in inpainting, and M is an orthogonal matrix that couples
the
RGB colormaps 3 .
M T is a matrix that de-couples the channels back to the original space.
In other words, one can view colorization as performing imputation in some spectral space.
Subsequently, for our colorization method we choose
[Equation]
Again, our forward measurement matrix is orthogonal, and we choose A such that we only affect the orthogonal complement of the measurement subspace.
The sampler for colorization is based on the predictor-corrector (PC sampler we choose to apply MCG y − C are chosen as hyper-parameters
CT Reconstruction For the case of CT reconstruction, the forward model reads
[Equation]
where R is the discretized Radon transfrom [7] that measures the projection images from different angles.
Note that for CT applications, R T corresponds to performing backprojection (BP), and R † corresponds to performing ﬁltered backprojection (FBP).
We choose
[Equation]
where the choice of A reﬂects that the Radon transform is not orthogonal, and we need the term ( RR T ) † as a term analogous to the ﬁltering step.
Indeed, this form of update is known as the algebraic reconstruction technique (ART), a classic technique in the context of CT [16].
We note that this choice is different from what was proposed in [40], where the authors repeatedly apply projection/FBP by explicitly replacing the sinogram in the measured locations.
From our experiments, we ﬁnd that repeated application of FBP is highly numerically unstable, often leading to overﬂow.
This is especially the case when we have limited resources for training data (we use 4k, whereas [40] uses 50k), as we further show in Section 5.
Algorithm for SV-CT reconstruction uses PC sampler (VE-SDE), where we use MCG step after one sweep of corrector-predictor update.
We note that this is a design choice, and one may as well use the MCG update step after both the predictor and corrector steps, as was proposed in [41].
We set
N = 2000 , α = 0 .
1 / (cid:107) R †
( y − R ˆ x 0 ) (cid:107) .
D Generative process of the proposed method
In Fig.
7, we depict the comparison of the generative process between the two methods: scoreSDE
[41], which relies on alternating projections; and our method, which utilizes MCG as correcting steps.
In Fig.
7 (a), we can clearly see the unnatural boundary between the masked and the unmasked region forming, and evolving as t → 0 , without getting corrected (Visible more clearly in ˆ x 0 ).
On the other hand, thanks to the additional gradient step that corrects the errors in the boundary, we see a much more natural evolution of the signal as t → 0 in Fig.
7 (b).
E Limitations
There exists a limitation speciﬁcally for the ImageNet dataset when using the proposed algorithm for inpainting.
Speciﬁcally, as shown in
Fig. 8, for the case of half-mask (i.e. the left or right half of the image is zeroed-out), we often see the reconstructions are generated showing symmetries that are unrealistic.
Note that this kind of effect is not observed in our FFHQ experiments.
Hence, we conjecture that this phenomenon arises from the imperfectness of the learned score function s θ .
Namely, due to the
ImageNet dataset being much more diverse and therefore widely known to be a much harder dataset to learn, the subobtimality of the score function may be greater than the FFHQ score function.
This could possibly lead to such deﬁciencies.
F Experimental Details F.1 Implementation details
Training of the score function For inpainting experiments, we take the pre-trained score functions that are available online (FFHQ 4 , imagenet 5 ).
For CT reconstruction experiment, we train a ncsnpp model with default conﬁgurations as guided in [41] with the VE-SDE framework.
The model was trained for 200 epochs with the full training dataset, with a single RTX 3090 GPU.
Training took about one week wall-clock time.
Required compute time for inference All our sampling steps detailed in Algorithm C was performed with a single RTX 3090 GPU.
The inpainting algorithm based on ADM [13] takes about 90 seconds (1000 NFE) to reconstruct a single image of size 256 × 256.
Our colorization and CT reconstruction algorithm based on score-SDE [41] takes about 600 seconds (4000 NFE) to infer a single 256 × 256 image.
Code Availability We will open-source our code used in our experiments upon publication to boost reproducibility.
F.2 Comparison methods
F.2.1 Inpainting,
Colorization
Score-SDE Score-SDE [41] demonstrated that unconditional diffusion models can be adopted to various inverse problems, such as inpainting and colorization.
Our method without the MCG step is identical to score-SDE, and hence we use the same score function, parameters, and sampler as used in the proposed method for reconstruction.
RePAINT RePAINT
[32] proposes to iterate between denoising-noising steps multiple times in order to better incorporate inter-dependency between the known and the unknown regions in the case of image inpainting.
We use the same score function and sampler for RePAINT as in the proposed method.
Following the default conﬁgurations in [32], we take N = 200 (corresponding to T in [32]), and U = 10 , where U denotes the count of iterated denoising-noising steps used within a single update index i .
4 https://github.com/jychoi118/ilvr_admhttps://github.com/openai/guided-diffusion
DDRM DDRM [25] demonstrates that linear inverse problems can be solved via diffusion models by decomposing the generative process with singular value decomposition (SVD), and performing reverse diffusion sampling in the spectral space.
The same score function adopted for the proposed method is used.
Using the notations from [25], we choose σ y = 0 , as we are aiming to solve noiseless inverse problem, and
η = 0 .
85 , η b = 1 .
The number of NFE is set to 20 with the DDRM sampling steps.
LaMa LaMA contains fast Fourier convolution in generator architecture for reconstructing images.
We trained the model from scratch using adversarial loss with r1 regularization term with its coefﬁcient 10 and gradient penalty coefﬁcient 0.001.
Adam optimizer is used with the ﬁxed learning rate of 0.001 and 0.0001 for discriminator network.
For FFHQ and Imagenet dataset, 500k iterations of trainings were done with batch size of 8.
AOT-GAN AOT-GAN consists of a deep image generator with a AOT block which consists of multiple length of residual blocks in parallel.
The discriminator is the same architecture with
PatchGAN from [55].
We trained the model from the scratch with 0.0001 learning rate using Adam optimizer β 1 = 0 and β 2 = 0 .
9 for both FFHQ and Imagent dataset.
500k iterations of trainings were done with batch size of 8.
Also, for style loss and the perceptual loss, VGG19 [39] pretrained on
ImageNet [12] was used.
ICT Image completion transformer (ICT) consists of two modules - a transformer model that follows the tokenization procedure to process information in the lower dimensional space, and another guided upsampling module to retrieve the data dimensionality.
The encoded features are sampled from a probability distribution via Gibbs sampling, such that one can capture multimodal reconstructions from the same measurement.
For both the FFHQ and Imagenet dataset, we used pretrained models provided by the authors.
IAGAN Image adaptive GAN (IAGAN) uses a pre-trained generator and adapts it at test time for the given forward model.
Speciﬁcally, following compressed sensing using generative model (CSGM) [4], one initializes the latent vector z such that z ∗ =
arg min z (cid:107) y − AG
θ ( z ) (cid:107) .
Then, the latent code and the neural network parameters are jointly optimized through some iterations of z ∗∗ , θ ∗ = arg min z ,θ (cid:107) y − AG
θ ( z ) (cid:107) .
The ﬁnal result is achieved by the forward pass through the generator, after which follows the projection into the measurement subspace.
For tuning the generator, we follow the default conﬁgurations from the ofﬁcial codebase.
Since the codebase uses a GAN that generates 1024 × 1024 images, we downscale the result into 256 × 256 image as a ﬁnal post-processing step.
DSI DSI is structured with the combination of VQ-VAE [46], structure generator and texture generator.
The architectures were trained separately, with Adam optimizer.
When inference, only structure and texture generator was used.
We trained the model from scratch.
During optimization, the structure generator used linear warm-up schedular and square-root decay schedule used in [36].
We used Adam optimizer on training all models with learning rate of 0.0001 and β 1 = 0 .
5 using exponetial moving average (EMA).
Training was done for 500k iteration for both FFHQ and Imagenet dataset.
cINN cINN is an invertible neural network which can take in additional conditions as input, and in our case grayscale images.
We train the model using default conﬁgurations as advised in https://github.com/VLL-HD/conditional_INNs without modiﬁcations.
FFHQ model was trained with the learning rate of 0.0001 for 100 epohcs using the Adam optimizer.
LSUN bedroom model was trained with the learning rate of 0.0001 for 30 epochs.
pix2pix Pix2pix is a variant of conditional GAN (cGAN) that takes in as input, the corrupted image.
The model is trained in a supervised fashion, with the loss consisting of the reconstruction loss, and the adversarial loss.
As the discriminator architecture, we adopt patchGAN [21], and utilize the
LSGAN [33] loss, weighting the adversarial loss by the value of 0.1.
Similar to cINN, FFHQ model was trained with the learning rate of 0.0001 for 100 epochs using Adam optimizer.
LSUN bedroom model was trained with the same conﬁguration for 30 epochs.
F.
CT
reconstruction
Score-CT
We use the hyper-parameters as advised in [40] and set η = 0 .
246 , λ = 0 .
The measurement consistency step is imposed after every corrector-predictor sweep as in the proposed method.
SIN-4c-PRN Directly using the ofﬁcial implementation 6 [50], we train the sinogram inpainting network (SIN) with the AAPM dataset for 200 epochs with the batch size of 8, and learning rate of 0.0001.
We train two models separately for different number of views - 18, and 30.
cGAN
We adopt the implementation of cGAN [15] from SIN-4c-PRN repository 6 .
We train the two separate networks for 18 view, and 30 view projection, with the same conﬁguration - 200 epochs, learning rate of 0.0001, and batch size of 8.
FISTA-TV We perform FISTA-TV [3] reconstruction using TomoBAR [45], together with the
CCPi regularization toolkit [27].
Leveraging the default setting, we use the least-squares (LS) data model, and run the FISTA iteration for 300 iterations per image, with the total variation regularization strength set to 0.001.
G Further Experimental Results
We provide extensive set of comparison study for each task in Fig.
9, 11, and 12.
Furthermore, in order to illustrate the ability of our method to generate multimodal reconstructions given a measurement, we present further experimental results of inpainting and colorization in the following ﬁgures: Fig. 13, 14, 15, and