Decomposed Diffusion Sampler for Accelerating Large-Scale Inverse Problems
Authors: Hyungjin Chung
Suhyeon Lee
Jong Chul Ye
Venue: ICLR 2024 poster (2024)
Keywords: Diffusion models; Inverse problems; Krylov subspace
Primary area: generative models

Abstract: Krylov subspace, which is generated by multiplying a given vector by the matrix of a linear transformation and its successive powers,  has been extensively studied in classical optimization literature to design algorithms that converge quickly for large linear inverse problems. For example, the conjugate gradient method (CG), one of the most popular Krylov subspace methods, is based on the idea of minimizing the residual error in the Krylov subspace. However, with the recent advancement of high-performance diffusion solvers for inverse problems, it is not clear how classical wisdom can be synergistically combined with modern diffusion models. In this study, we propose a novel and efficient diffusion sampling strategy that synergistically combines the diffusion sampling and Krylov subspace methods. Specifically, we prove that if the tangent space at a denoised sample by Tweedie's formula forms a  Krylov subspace, then the CG initialized with the denoised data ensures the data consistency update to remain in the tangent space. This negates the need to compute the manifold-constrained gradient (MCG), leading to a more efficient diffusion sampling method. Our method is applicable regardless of the parametrization and setting (i.e., VE, VP). Notably, we achieve state-of-the-art reconstruction quality on challenging real-world medical inverse imaging problems, including multi-coil MRI reconstruction and 3D CT reconstruction. Moreover, our proposed method achieves more than 80 times faster inference time than the previous state-of-the-art method. Code is available at https://github.com/HJ-harry/DDS

Published as a conference paper at ICLR 2024
DECOMPOSED DIFFUSION SAMPLER FOR ACCELERAT- ING LARGE-SCALE INVERSE PROBLEMS
Hyungjin Chung', Suhyeon , Jong Chul Ye²
1 Dept. of Bio & Brain Engineering, KAIST, 2 Kim Jae Chul Graduate School of AI, KAIST
(hj chung, Jdhyeon.lee, jong.yez yeg..yeeekiise ac.
ABSTRACT
Krylov subspace, which is generated by multiplying a given vector by the matrix linear transformation and its successive powers, has been extensively studied in classical optimization literature to design algorithms that converge quickly for large linear inverse problems. For example, the conjugate gradient method (CG), one of the most popular Krylov subspace methods, is based on the idea of minimizing the residual error in the Krylov subspace. However, with the recent advancement of high-performance diffusion solvers for inverse problems, it is not clear how classical wisdom can be synergistically combined with modern diffusion models. In this study, we propose novel and efficient diffusion sampling strategy that synergistically combines the diffusion sampling and Krylov subspace methods. Specifically, prove that if the tangent space denoised sample by Tweedie's s formula forms a Krylov subspace, then the CG initialized with the denoised data ensures the data consistency update remain in the tangent space. This negates the need to compute the manifold-constent efficient diffusion sampling method. Our method is applicable regardless of the parametrization and setting (i.c., VE, VP). Notably, we achieve state-of-the-art reconstruction quality on challenging real-world medical inverse imaging problems, including multi-coil MRI reconstruction and 3D CT reconstruction. Moreover, our proposed method achieves more than 80 times faster inference time than the previous state-of-the-art method. Code is available /tltttps://gihhb.. com/ HJ-harry/DDS
gradient (MCG), leading to a more
INTRODUCTION
Diffusion models (Ho et al. 2020 Song et al. 2021b) are the state-of-the-art generative model that learns to generate data by gradual denoising, starting from the reference distribution (e.g. Gaussian). In addition to superior sample quality in the unconditional sampling of the data distribution, it was shown that one can use unconditional diffusion models as generative priors that model the data distribution (Chung et al., 2022a Kawar et al. 2022), and incorporate the information from the forward physics model along with the measurement y to sample from the posterior distribution p(xly). This property is especially intriguing when seen in the context of Bayesian inference, can separate the parameterized prior (ar) from the measurement model (i.e. likelihood) p(y a:) to the same pre-trained
construct the posterior po(x)y) po( (r)p((y)) ar). In other words, neural network model regardless of the forward model at hand. Throughout the manuscript, we refer this class of methods as Diffusion model-based Inverse problem Solvers (DIS).
For inverse problems in medical imaging, e.g. magnetic resonance imaging (MRI), computed tomography (CT), is often required accelerate the measurement process by reducing the number of measurements. However, the data acquisition scheme may vary vastly according to the circumstances (e.g. vendor, sequence, etc.), and hence the reconstruction algorithm needs be adaptable to different possibilities. Supervised learning schemes show weakness in this aspect as they overfit to the measurement types that were used during training. As such, is easy to see that diffusion models will be particularly strong in this aspect as they are agnostic to the different forward models at inference. Indeed, it was shown in some of the pioneering works that diffusion-basee reconstruction
algorithms have high generalizability (Jalal et al. 2021 Chung & Ye 2022 Song al. 2022 Chung
et al. 2023b).

---

Published as a conference paper at ICLR 2024
(4) Multi-ool CS-MIU
Meck
SON FOOCE
EUAN
Nttelles
31.6-00 IUSAI
S2.12/ 997
33, 30, FLRS7
34 64 OURRS
Nikey = 0.05)
MSS DIXK
(h) 3D Sparse-view CT
1DP
Di (400)
Ourss (35)
Canond Tuuhh
3296 0R32
34, NANNUR
36. St. one
Ouas (19)
Figure 2: Representative reconstruction results. (a) Multi-coil MRI reconstruction, (b) 3D sparse-view CT. Numbers parenthesis: NFE. Yellow numbers in bottom left corner: PSNR/SSIM.
While showing superiority in performance and generalization capacity, slow inference time is known to be critical drawback of diffusion models. One of the most widely acknowl- edged accelerated diffusion sampling strate- gies is the denoising diffusion implicit model (DDIM) (Song et al. 2021a), where the stochas- tie ancestral sampling of denoising diffusion probabilistic models (DDPM) can be transi- tioned to deterministic sampling, and thereby accelerate the sampling. Accordingly, DDIM sampling has been well incorporated solving low-level vision inverse problems (Kawar et al. 2022, Song et al. 2023). In recent applica- tion of DDIM for linear image restoration tasks, Wang et al. 2023) proposed an algorithm dubbed denoising diffusion null-space models (DDNM), where one-step null-space modification is made to impose consistency. However, the sampling strategy is not successful in practical large-scale medical imaging contexts when the forward model significantly more complex (e.g. parallel imaging (PI) CS-MRI, 3D modalities). Furthermore, is unclear how the algorithm related to the existing literature of conditional sampling approaches that take into account the geometry of the manifold (Chung et al.) 2022a 2023a).
25N0180
30.0
50.0
DBS (Am
Figure 1: Parallel imaging MR reconstruction eval- uation PSNR vs. NFE (log scale). Reconstruction from ID uniform random x4 acceleration (Zbon- tar al. 2018).
On the other hand, in classical optimization literature, Krylov subspace methods have been extensively studied to deal with large-scale inverse problems due their rapid convergence rates (Liesen &
2

---

Published as a conference paper at ICLR 2024
Strakos 2013). Specifically, consider a typical linear inverse problem
Ax,
where is the linear mapping and the goal is to retrieve x from the measurement y. Without loss of generality, throughout the paper we assume that A in ① is square. Otherwise, we can obtain an equivalent inverse problem with symmetric linear mapping A as ÿ := A'y A²Ax == Az. This is because the solution to the normal equation A Ax A"y y is indeed a solution to Ax = y if A has full column rank, which holds in most of the ill-posed inverse problem cases. Given an initial guess ăt Krylov subspace methods seek an approximate solution (! from affine subspace + Kr, where the 1-th order Krylov subspace K, defined by
K2 == (b, Ab, ... A)) b: y= Aá
(2)
For example, the conjugate gradient (CG) method is special class of the Krylov subspace method that minimizes the residual in the Krylov subspace. Krylov subspace methods are particularly useful for large-scale problems thanks to their fast convergence (Liesen & Strakos 2013).
Inspired by this, here we are interested in developing method that synergistically combines Krylov subspace methods with diffusion models such that can be effectively used for large-scale inverse problems. Specifically, based on novel observation that diffusion posterior sampling (DPS) (Chung) et al. 2023a) with the manifold constrained gradient (MCG) (Chung et al., 2022a) is equivalent to one-step projected gradient the tangent space at the "denoised" data by Tweedie's 's formula, provide multi-step update scheme on the tangent space using Krylov subspace methods. Specifically, show that the multiple CG updates are guaranteed to remain in the tangent space, and subsequently generated sample with the addition of the noise component can be correctly transferred to the next noisy manifold. This eliminates the need for the computationally demanding MCG while permitting multiple economical CG steps at each ancestral diffusion sampling, resulting in a more efficient DDIM sampling. Our analysis holds for both variance- preseervn (VP) and variance-expea -expooiie (VE) sampling schemes.
The combined strategy, dubbed Decomposed Diffusion Sampling (DDS), yields better performance with much-reduced sampling time (20~50 NFE, x80 ~ 200 acceleration; See Fig. for comparison, for representative results), and is shown to be applicable to a variety of challenging large scale inverse problem tasks: multi-coil MRI reconstruction and 3D CT reconstruction.
2 BACKGROUND
Krylov subspace methods Consider the linear system 11. In classical projection based methods such as Krylov subspace methods (Liesen & Strakos 2013), for given two subspace K. and C, we define approximate problem to find æ K such that
- Ax L
This is basic projection step, and the sequence such steps is applied. For example, with non-zero estimate é, the associated problem is to find E & + K. such that y Ax 1 L, which is equivalent to finding K. such that
A6 1. C, 8 x.
y Aá.
In terms of the choice of the two subspaces, the CG method chooses the two subspaces IC and L as the same Krylov subspace:
K= Abb, ....
Ab.....AA-
(5)
Then, CG attempts to find the solution to the following optimization problem:
min
1
|
(6)
#Ea+K:
Krylov subspace methods can be also extended to nonlinear problems via zero-finding. Specifically, the optimization problem ming {(x) can be equivalently converted to a zero-finding problem of its gradient, i.c. V v_[((r) {(x) = 0. If we consider a non-linear forward imaging operator A(-). we can define ((x)) 1y A(x) || 2 /2. Then, one use, for example, Newton-Krylon method (Knoll & Keyes
3

---

Published as a conference paper at ICLR 2024
2004) to linearize the problem near the current solution and apply standard Krylov methods to solve the current problem. Now, given the optimization problem, we can see the fundamental differences between the gradient-based approaches and Krylov methods. Specifically, gradient methods are based on the iteration:
x: (1+1) (t) (al((x)))
which stops updating when ))) 0. On the other hand, Krylov subspace methods try find E K2 by increasing / to achieve better approximation of (x)(x)= = 0. This difference allows us devise a computationally efficient algorithm when combined with diffusion models, which we investigate in this paper. See Appendix A.2 fur further mathematical background.
Diffusion models Diffusion models (Ho et al., 2020) attempt to model the data distribution Pdaat (xz) by constructing a hierarchical latent variable model
T
po(xo) = =f (o(T)) II
(2) (2)-1120)) dæ
(8)
t=1
where T) E Rd are noisy latent variables that have the same dimension as the data random vector 2o Rd, defined by the Markovian forward conditional densities
(<<,|<1-1) = (1 B.) g (V. v0, (1 T.)
(9)
(10)
Here, the noise schedule 3e an increasing sequence of I, with ār : II or, a == 1 - 3r. Training of diffusion models amounts to training multi-noise level residual denoiser (i.e. epsilon matching)
min E Fe qam (F0)
[1c55] (xx) (1,)
c1B]
such that cg- (xx)= 22 /a
Furthermore, it can be shown that epsilon matching is equivalent to Ermon 2019) objective up
a az
the denoising score matching (DSM) (Hyvärinen Dayan 2005 Song
to constant with different parameterization
min Ez, Bo.e | (1) (xx) - V Ve, log (xx)((10))
(11)
such that sp. (xx) 22- /0,20
(1) (xx)) //I 1 ā. For notational simplicity, we often denote cg.
1-a
81, ét, de instead of s (ar)), 60 (xxe ), Dz- (zx). representing the estimated score, noise, and noiseless
image, respectively. Sampling from (8) can be implemented by ancestral sampling, which iteratively performs
- at és ) + Bie,
T1-1
xt
(12)
Vai
V1 αt
where 3, :
1-6, I-as 32.
One can also view ancestral sampling 122 as solving the reverse generative stochastic differential equation (SDE) of the variance preserving (VP) linear forward SDE (Song et al. 2021b). Additionally, one can construct the variance exploding (VE) SDE by setting (< = A (a, Zn, of I), which
in form of Brownian motion. In Appendix A.I models under the score/SDE perspective.
we further review the background on diffusion
Seen either from the variational or the SDE perspective, diffusion models are inevitably slow sample from. To overcome this issue, DDIM (Song et al., (021a) proposes another method of sampling which only requires matching the marginal distributions q( @. wo). Specifically, the update rule is given as follows
DDIM
Tt-1 = Vot-12 + V ā.-1 - 1/
(1) (a.) + nBEE,
(13)
= Vār-iē + w;;
4

---

Published as a conference paper at ICLR 2024
where 2, is the denoised estimate
(1) zz Zz- (xt)
_(x; v1 ā.c (xx)).
(14)
Vā
which can also be equivalently derived from Tweedie's formula Efron (201 I), and w, denotes the total noise given by
We v 1 (4) (ax;) ) + 1)BrE
(15)
In 13 13.. E [0, 1] parameter controlling the stochasticity of the update rule: 1) = 0.0 leads to fully
deterministic sampling, whereas 1) = 1.0 with 3. = v(1 â. 1)((1 (4)) V 0x/0 recovers the ancestral sampling DDPMs.
It is important to note that the noise component w. properly matches the forward marginal (Song al. 2021a). The direction w, of this transition is determined by the vector sum of the deterministic and the stochastic directional component. Accordingly, assuming optimality of Eg-, the total noise W in 15 can be represented by
(16) for some ~ (o, 1) (see Appendix B In other words, 13 is equivalently represented by 2t-1 10 VI 14-16 for some (O, I). Therefore, it can be seen that the difference between DDIM and DDPM lies only in the degree of dependence on the deterministic estimate of the noise component with feasible intermediate values 7) E (0, 1).
W0
DECOMPOSED DIFFUSION SAMPLING
Conditional diffusion for inverse problems
The conditional diffusion sampling for inverse prob-
lems (Kwarr et al. 2022 Chung et al., 2023a b) attempts to solve the following optimization problem:
min {(x)
(17)
#cM
where {(x) denotes the data consistency (DC) loss (i.e., f(x) = by - Axl12/2 for linear inverse problems) and M represents the clean data manifold. Consequently, it is essential to navigate in way that minimizes cost while also identifying the correct clean manifold. Accordingly, most of the approaches use standard reverse diffusion (e.g. (2p), alternated with operation to minimize the DC loss.
Recently, Chung al. (2023a) proposed DPS, where the updated estimate from the noisy sample 2t E .M, is constrained to stay on the same noisy manifold M, This is achieved by computing the MCG (Chung al. 2022a) on noisy sample æ; M, Vmeg f(x;) == Va, ((ē;). where &, is the denoised sample in (14 through Tweedie's 's formula. The resulting algorithm can be stated as follows:
Tt-1 v (#. V ((x))) + w;
(18)
where "/e > 0 denotes the step size. Since parameterized score function <g- (ax) is trained with samples supported on Mr, e) shows good performance on denoising I; M,, allowing precise transition to MM-1- Therefore, by performing
(88) from / = T to f O. we can solve the optimization problem (17) with 2n E M. Unfortunately, the computation of MCG for DPS requires
computationally expensive backpropagation and is often unstable (Poole et al. 2022 Du et al. (023).
Key observation By applying the chain rule for the MCG term in we have
Fd, Da.
Va, ((t.)
v ((xe)
where we use the denominator layout for vector calculus. Since V ((x)) is standard gradient, the main complexity of the MCG arises from the Jacobian term Bm:
In the following Proposition]] | we show that if the underlying clean manifold forms affine subspace, then the Jacobian term dæ:
diz is indeed the orthogonal projection the clean manifold up scale
5

---

Published as a conference paper at ICLR 2024
factor. Note that the affine subspace assumption is widely used in the diffusion literature that has been used when 1) studying the possibility of score estimation and distribution recovery (Chen et al.) 2023), 2) showing the possibility of signal recovery (Rout et al. 2023a b). and the most relevantly, 3) showing the geometrical view of the clean and the noisy manifolds (Chung et al. (022a)). Although it is difficult to assume in practice that the clean manifold forms affine subspace, it could be approximated by piece-wise linear regions represented by the tangent subspace at ă, Therefore, Proposition
D is still valid in that approximate regime.
Proposition (Manifold Constrained Gradient). Suppose the clean data manifold M is represented as an affine subspace and assumes the uniform distribution on M. Then,
1
PM
(19)
dz.
Vā,
7/ V m. ((at))
PM (#t Ca V. ((tre)
(20)
for some > 0, where PMM denotes the orthogonal projection to M.
Now, 20 in Proposition I [ndiiaaes that if the clean data manifold is an affine subspace, the DPS corresponds the projected gradient on the clean manifold. Nonetheless, a notable limitation of MCG is its inefficient use of a single projected gradient step for each ancestral diffusion sampling. Motivated by this, we aim to explore extensions that allow computationally efficient multi-step optimization steps per each ancestral sampling.
Specifically, let T. denote the tangent space on the clean manifold at a denoised sample é; 14 Suppose, furthermore, that there exists the i-th order Krylov subspace:
K1,1 Sban(b, Ab... ... A' A'-Ib), b b: y Aa
(21)
such that
T. u. + Kr,1
Then, using the property of CG in (6). it is easy to see that M M-step CG update with M < starting from & are confined in T7 since corresponds to the solution of
min mER; K.M
1y Axll
(22)
and KM C KG when M < 1. This offers a pivotal insight. It shows that if the tangent space at each denoised sample is representable by a Krylov subspace, there's no need to compute the MCG. Rather, the standard CG method suffices to guarantee that the updated samples stay within the tangent space. To sum up, our DDS algorithm is as follows:
Pt-1 = Vou- a + Wr.
(23)
d// = CG(A A. A C., M),
M
(24)
where CG(-) denotes the M-step CG for the normal equation starting from ã.. In contrast, DDNM (Wang al. 2023) for noiseless image restoration problems uses the following update instead of (24):
a/ (I AtA)a, + A'y,
(25)
where A denotes the pseudo-inverse of A. Unfortunately, (55 in DDNM does not ensure that the update signal é;, lies in Te due to A
Therefore, for large-scale inverse problems, we find that CG outperforms naive projections (25) by quite large margins. This is to be expected, as CG iteration enforces the update to stay on T, whereas the orthogonal projections in DDNM do not guarantee this property. In practice, even when our Krylov subspace assumptions cannot be guaranteed, we empirically validate in Appendix F. that DDS indeed keeps æ; closest to the noisy manifold Mr. which, in turn, shows that DDS keeps the update close to the clean manifold M. Moreover, it is worth emphasizing that maddeent-based methods (Jalal et al. 2021 Chung et al. 2022a 2023a) often fail when choosing the "theoretically correct" step sizes of the likelihood. To fix this, several heuristics on the choice of step sizes required (c.g. choosing the step size x 1/1ly A&, 2|2), which easily breaks when varying the NFE. In this regard, DDS is beneficial in that is free from the cumbersome step-size tuning process.

---

Published as a conference paper at ICLR 2024
Suankall
N014
Uniform ID
x4
PANR (ub)
SSIM
27 59 =8:43
0667:0117
31,77u0
0380c 11
0.5660011
33,8802 10
0904 08
0:877:0.0
33.25c1
0180003 mg
30.66-100
0896-2111
0,856 0M
36A1:992
0107700 18
PSNR (ub]
25
31.5800.35
69
33,01-234
30.37
X PSKR joh)
1746-00 m6
M01563.00 9
11.441 95
14.55-1 56
11.47:1
15 12
35.15 9:20
Ganasiam ID
S50M
0.709-005
0365
08781663
0 14
0.835-0000
0.81
0.451 0.00
0.997
PSNR (ib]
124
11
32.5622
32.43-9156
30,472230
3548-900
32,87
SSIM
0.4723822
08412010
0833p0 15
0
0,3922.18
093732417
0.947:9.15
0.94003311
PSNR (ch)
24
31V711:34
3:
51.41-133
20.5
33.99:3:301
12.95 56
Canssian
PSKR (sh)
1491-308
16 06666.399
27 34:4 05
28:17:000
76, 30.00.1000 31
37 06:66-110 11
75. 1-275
25
SSIM
0.54738.33
0,32200.29
045531429
0.0920135 15
0,704
0.65533111
0.732:2:10
0 0.695
0.692
VD Paisonn disk
8
SSIM
0.56218:1
0.375 10
0.9921033
0833.010
0.865:1177
08015010
0.99215917
0.815..31
0.915
PSNR (oh]
19:5:35
16:1-:500
18:213.000
90.72:1 50
20:99:1
1002-11
14.99-1
SSIM
0.
0.54821229
0.749
0,202
0.753
0034:3/0
0.931
Table 1: Quantitative metrics for noiseless parallel imaging reconstruction. Numbers in parenthesis:
NFE. ": expressed as x2 x as the network is evaluated for real/imag part of each coil. Bold: best.
Mean + 1 std.
Furthermore, our CG step can be easily extended for noisy image restoration problems. Unlike the DDNM approach that relies on the singular value decomposition to handle noise, which is non-trivia to perform on forward operators in medical imaging (e.g. PI CS-MRI, CT), we can simply minimize the cost function
((x)) - - Axl2
2
+ 2 || - az 12
(26)
by performing CG iteration CG(> + I. + 7''y, at. M) in the place of (24). where 7 is a aaper-parameter that weights the proximal regularization (Parikh & Boyd| 2014). Finally, our method can also be readily extended to accelerate DiffusionMBIIR (Chung et al. 2023b) for 3D CT reconstruction by adhering to the same principles. Specifically, we implement an optimization strategy impose the conditioning:
min 2 Aa 31 118++IDD D|
(27)
where D= is the finite difference operator that is applied to the z-axis that is not learned through the diffusion prior, and unlike Chung et al. (2023b), the optimization is performed in the clean manifold starting from the denoised i; rather than the noisy manifold starting from ææ. As the additional prior only imposed in the direction that is orthogonal to the axial slice dimension (ry) captured by the diffusion prior (i.e. manifold .M). (7) can be solved effectively with the alternating direction method of multipliers (ADMM) (Boyd et al. [011]) after sampling 2D diffusion slice by slice. See Appendix for details in implementation.
4 EXPERIMENTS
Problem setting. We have the following general measurement model (see Fig. 3 for illustration of the imaging physics).
= PTSæ =: Ax, y E C", E Caxd
(28)
where P is the sub-sampling matrix, T is the discrete transform matrix (i.e. Fourier, Radon), and when we have single-array measurement including CT, and [1(1) ,..., s(() when we have a coioll parallel imaging (PI) measurement.
We conduct experiments on two distinguished applications applications-accllereea accelerated MRI. and 3D CT recon- struction. For the former, we follow the evaluation protocol of Chung & Ye (202)) and test our method on the fastMRI knee dataset (Zhontar et al. 2018) on diverse sub-sampling patterns. We provide comparisons against representative DIS methods Score-MRI (Chung & Ye 2022) Jalal et al., (2021). DPS (Chung et al.) 2023a). Notably, for DPS, we use the DDIM sampling strategy to show that the strength of DDS not only comes from the DDIM sampling strategy but also the use of the sampling together with the CG update steps. The optimal 7) values for DPS (DDIM) are obtained through grid search. We do not compare against (Song et al. 2022) the method cannot cover the multi-coil setting. Other than DIS, we also compare against strong supervised learning baselines: E2F- Varnet (Sriram al. 2020), U-Net (Zbontar et al. 2018); and compressed sensing baseline: total
7

---

Published as a conference paper at ICLR 2024
variation reconstruction (Block et al. 2007). ). For the latter, we follow Chung al. (2023b and test both vparse-view CT reconstruction (SV-CT), and limited angle CT reconstruction (LA-CT) on the
AAPM 256 x 256 dataset. We compare against against representative DIS methods DifIusionMBIR (Chung
et al. 2023b), Song et al. (2022). MCG (Chung et al. 2022a), and DPS (Chung et al. 2023a):
Supervised baselines Lahiri et al. (2022) and FBPConv Net (Jin et al. 2017); Compressed sensing baseline ADMM-TV. For all proposed methods. we employ M = 5, 0.15 for 19 NFE, 0.5
for 49 NFE, 77 0.8 for 99 NFE unless specified otherwise. While we only compare against single CS baseline. it was reported in previous works that diffusion model-based solvers largely outperform the classic CS baselines (Jalal et al. 2021 Luo et al., 2023). for example, LI-wavelet (Lustig et al.) 2007) and LI-ESPiRiT (Uecker et al. 2014). For PI CS-MRI experiments, we employ the rejection sampling based on a residual-based criterion to ensure stability. Further experimental details can be found in appendix G
4.1 ACCELERATED MRI
Improvement over DDNM. Fixing the sampling strategy the same, inspect the effect of the three different data consistency imposing strategies: Score-MRI (Chung & Yel 2022). DDNM (Wang) al.) (023), and DDS. For DDS, we additionally search for the optimal number of CG iterations per sampling step. In Tab. 2 we see that under the low NFE regime, the score-MRI DC strategy has significantly worse performance than the proposed methods, even when using the same DDIM sampling strategy. Moreover, we see that overall, DDS outperforms DDNM by a few db in PSNR. We see that 5 CG iterations per denoising step strike a good balance. One might question the additional computational overhead of introducing the iterative CG into the already slow diffusion sampling. Nonetheless, from our experiments, we see that on average, a single CG iteration takes about 0.004 sec. Consequently, it only takes about 0.2 see more than the analytic counterpart when using 50 NFE (Analytic: 4.51 see vs. CG((): 4.71 sec.).
VE (Chung & Ye 2022). Keeping the pre-trained model intact from Chung & Ye (2022), we switch from the Score-MRI sampling to Algorithm and report on the reconstruction results from uniform ID x4 accelerated measurements in Tab. 6 Note that Score-MRI uses 2000 PC as the default setting, which amounts to 4000 NFF, reaching 33.25 PSNR. We see almost no degradation in quality down to 200 NFE, but the performance rapidly degrades as we move down to 100, and completely fails when we set the NFE < 50. On the other hand, by switching to the proposed solver, are able to achieve the reconstruction quality that better than Score-MRI (4000 NFE) with only 100 NFE sampling. that we can reduce the NFE down to 30 and still achieve decent reconstructions. This is a useful with speed. However, we observe several downsides of using the VE parameterization including numerical instability with large NFE, which we analyze in detail in appendix
Improvement
DDS (ours)
Score-MRI DDNM
3
5
10
PSNR(db] SSIM
26.48
31.51 33.78 34.61 32.48 0,932
31.36
0.688
0.934
0.952 0.956 0.949
Table 2: Ablation study on the DC strategy. 49 NFE VP DDIM sampling strategy, uniform ID x4 acc. reconstruction.
Moreover, property for a reconstruction algorithm, as we can trade off reconstruction quality
Parallel Imaging with VP parameterization (Noiseless). We conduct thorough PI reconstruction experiments with 4 different types of sub-sampling patterns following Chung & Ye (2022). Algo- rithm 2 in supplementary material is used for all experiments. Quantitative results are shown in Tab. AAsso see Fig. 7 for qualitative results). As the proposed method is based on diffusion models, it is agnostic to the sub- saomsmmmling patterns, generalizing well all the different sampling patterns, whereas supervised learning-based
methods such as U-Net and Varnet fail dramatically on 2D subsampling patterns. Furthermore, to emphasize that the proposed method is agnostic to the imaging forward model, we show for the first time in the DIS literature that DDS is capable of reconstructing from con-cartesian MRI sub-sampling patterns that involve non-uniform Fast-Fourier Transform (NUFFT) (Fessler & Sutton 2003). See Appendix F2
In Tab. I we see that DDS sets the new state-of-the-ar in most cases even when the NFE is constrained to < 100. Note that this is dramatic improvement over the previous method Chung & Ye (2022). as for parallel imaging, Score-MRR required 120k(C = 15) NFE for the reconstruction of single image. Contrarily, DDS is able to outperform score-MRI with 49 NFE, and performs
8

---

Published as a conference paper at ICLR 2024
B-skew
4nkee
Setten
SagRoai
Sagitall
Methed
PSNNE F 5016
1SNN
35134
PSNIL
55034
73845 1 533M1
PSNNX
.SIIN
USNN
PSNNE
55054
FSNNL
SSUM
PSNNR
00381
7186 VP (19)
X3.66
0.930
15.30 14
0.034
13.47
63817
31.48
0.918
39.81
01919
11
11:30
25.85
8.877
34.60
4.871
23.69
0.791
D0S (99)
55.43
0/322
24.35
0072
0033
31:23
0.912
0.999
0.924
23.09
0.838
0.505
22.26
DilosseMRRR Mare d 1o0cl (4000)
53.49
0.342
35.18
0967
12.18
0910
30.53
0.914
30.09
0.31
27.59
0.871
34.11
0.210
23.15
0.641
21:22
0.786
DPS
Chorpat pring kn
27.86
0.854
27,07
0.960
13.65
0.784
26 96
0.542
36.09
0117
22.10
0.33
23.06
0.773
21.56
0.784
9.34
0.666
MCG
(0000)
2861
0.333
26,05
0.944
34.55
0.765
27.33
0.855
36.53
0.85
23.44
0.745
34.99
0:621
23.52
0.806
20.71
0.685
pardtal 2211
16.37
03333
10.13
0.714
18.11
6714
19.45
0.539
19.43
0.33
1544
06530
15.31
0:211
1340
0.531
11.07
0.483
ADAMM-TV
16.9
0635
1645
0.77
11.0
0716
13.99
0.688
1573
0.883
06638
16.36
0.00
13.77
0.666
1149
0.983
Table 4: Quantitative evaluation of SV-CT on the AAPM 256x256 test (mean values; std values
in Tab. 8). (Numbers in parenthesis): NFF, Bold: best.
on par with score-MRI with 19 NFE. Even disregarding the additional x2C more NFEs required for score-MRI to account for the multi-coil complex valued acquisition, the proposed method still achieves x80 ~ x x200 acceleration. We note that average, our method takes about 4.7 seconds for 49 NFE, and about 2.25 seconds for 19 NFE on a single commodity GPU (RTX 3090).
Noisy multi-coil MRI reconstruction. One of the most intriguing properties of the proposed DDS is the ease of handling measurement noise without careful computation of singular value decomposition (SVD), which non-trivial to perform for our tasks. With (26), we can solve
Mask Patlern
Acc.
TV
DPS (1000)
DDS (49)
Uniform ID
x44
PSNR [db]
SSIM
24.19
0.687
24.40
0.656
0.866
29,47
PSNR [db) 23.02
24.60
26.77
0.638
0.666
0.827
PSNR [Ubb]
23.07
23,48
3I,95
with CG, arriving at Algorithm 3 in supple- mentary material. For comparison, methods that try to cope with measurement noise via SVD in the diffusion model context (Kawar et al.] [2022:
VD Poisson disk
SSIM 0.609
0.592
0.890
15
PSNR [dh) 20.92
SSIM
0.554
0,622
23.57
0,853
24.36
Table 3: Quantitative metrics for noisy parallel imaging reconstruction. Numbers in parenthesis: NFE.
Wang al., (023) are not applicable and cannot be compared. One work that does not require computation of SVD and hence is applicable is
DPS (Chung et al. 2023a) relying backpropagation. To test the efficacy of DDS on noisy inverse problems, we add a rather heavy complex Gaussian noise (a = 0.05) to the k-space multi-coil measurements and reconstruct with Algorithm by setting y 0.95 found through grid search. In Tab. 3
we see that DDS far outperforms DPS (Chung et al. 2023a) with 1000 NFE by a large margin, while being about 40 faster as DPS requires the heavy computation of backpropagation.
RNCONSTRUCTION
4.2 3D CT RECONSTRUCTION
Sparse-view CT. Similar to the accelerated MRI experiments, we aim to both 1) improve the original VE model Chung et al. (2023b), and 2) train a new VP model better suitable for DDS. Inspecting Tab. 4 we see that by using Algorithm 6 in supplementary material, we are able to reduce the NFE to 100 and still achieve results that are on par with DiffusionMBIIR with 4000 NFF. However, we observe similar instabilities with the VE parameterization. Additionally, we find it crucial to initialize the optimization process with CG and later switch to ADMM-TV using a CG solver for proper convergence (see appendix E.2 for discussion). Switching to the VP parameterization and using Algorithm[4] 4 we now see that DDS achieves the new state-of- sfate-of-the-ans the-the-ar art with I 49 NFE. Notably, this
decreases the sampling time to 25 min for 49 NFE, and 10 min wall-clock time for 19 NFE single RTX 3090 GPU, compared to the painful 2 days for DiffusionMBIR. In Tab. we see similar improvements that were seen from SV-CT, where DDS significantly outperforms DiffusionMBIR while being several orders of magnitude faster.
CONCLUSION
In this work, we present Decomposed Diffusion Sampling (DDS), a general DIS for challenging real-world medical imaging inverse problems. Leveraging the geometric view of diffusion models and the property of the CG solvers on the tangent space, we show that performing numerical optimization schemes on the denoised representation is superior to the previous methods of imposing DC. Further, we devise fast sampler based on DDIM that works well for both VE/VP settings. With extensive experiments on multi-coil MRI reconstruction and 3D CT reconstruction, we show that DDS achieves superior quality while being 2 x80 faster than the previous DIS.
9

---

Published as a conference paper at ICLR 2024
Ethics Statement We recognize the profound potential of our approach to revolutionize diagnostic procedures, enhance patient care, and reduce the need for invasive techniques. However, we are also acutely aware of the ethical considerations surrounding patient data privacy and the potential for misinterpretation of generated images. All medical data used in our experiments were publicly available and fully anonymized, ensuring the utmost respect for patient confidentiality. We advocate for rigorous validation and clinical collaboration before any real-world application of our findings, to ensure both the safety and efficacy of our proposed methods in medical setting.
Reproducibility Statement For every different application and different circumstances (noise- less/noisy, VE/VP, 2D/3D), we provide tailored algorithns(See Appendix. ([E) to ensure maximum reproducibility. All the hyper-parameters used in the algorithms are detailed in Section 4 and Appendix G Code open-sourced athttpss://gthhu com/HJ--arrry/DSS
ACKNOWLEDGMENTS
This research was supported by the National Research Foundation of Korea(NRF)(RS-20)) 2023- 00262527), Field-oriented Technology Development Project for Customs Administration funded by the Korean government (the Ministry of Science & ICT and the Korea Customs Service) through the National Research Foundation (NRF) of Korea under Grant NRF202 I M331 & NRF- 2021M3I1 A 1097938, Korea Medical Device Development Fund grant funded by the Korea gov- ermment (the Ministry of Science and ICT, the Ministry of Trade, Industry, and Energy, the Min- istry of Health & Welfare, the Ministry of Food and Drug Safety) (Project Number: 1711 1137899, KMDF_PP PR_ 202000020 (0115), and Culture, Sports, and Tourism R&D Program through the Korea Creative Content Agency grant funded by the Ministry of Culture, Sports and Tourism 2023.
REFERENCES
Kai Tobias Block, Martin Uecker, and Jens Frahm. Undersampled radial MRI with multiple coils. Iterative image reconstruction using a total variation constraint. Magnetic Resonance in Medicine: An Official Journal of the International Society for Magnetic Resonance in Medicine, 57(6): 1086-1098, 2007.
Stephen Boyd, Neal Parikh, and Eric Chu. Distributed optimization and statistical learning via the alternating direction method of multipliers. Now Publishers Inc, 2011.
Guangyong Chen, Fengyuan Zhu, and Pheng Ann Heng. An efficient statistical method for image noise level estimation. In Proceedings of the IEEE International Conference on Computer Vision, pp. 477-485, 2015.
Minshuo Chen, Kaixuan Huang, Tuo Zhao, and Mengdi Wang. Score approximation, estima- tion and distribution recovery of diffusion models on low-dimensional data. arXiv preprint arXiv:2302.07194 2023.
Hyungjin Chung and Jong Chul Ye. Score-hased diffusion models for accelerated mri. Medical Image Analysis, pp. 102479, 2022.
Hyungjin Chung, Byeongsu Sim, Dohoon Ryu, and Jong Chul Ye. Improving diffusion models for inverse problems using manifold constraints. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022a. URL https://openreviiw. net met//forumiddnnJjjy1D111
Hyungjin Chung, Byeongsu Sim, and Jong Chul Ye. Sals-Cooooo-ffuue Conditional Diffusion Models for Inverse Problems through Stochastic Contraction. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2022b.
Accelerating
Hyungjin Chung, Jeongsol Kim, Michael Thompson Mecann, Mare Louis Klasky, and Jong Chul Ye. Diffusson posterior sampling for general noisy inverse problems. In International Conference on Learning Representations, 2023a. URL http::httttthteeee http//ppemeeiie..ete/oooorom
wwwp:ww.eeeeeee.... http///rrummm1000002222GGG 1==nn99GAGT@M
Hyungjin Chung, Dohoon Ryu, Michael T Mecann, Mare L Klasky, and Jong Chul Ye. Solving 3d inverse problems using trentrd 2d diffusion models. IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2023b.
10

---

Published as a conference paper at ICLR 2024
Prafulla Dhariwal and Alexander Quinn Nichol. Diffusion models beat GANs on image synthesis. In A. Beygelzimer, Y. Dauphin, Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021.
Yilun Du, Conor Durkan, Robin Strudel, Joshua B Tenenbaum, Sander Dieleman, Rob Fergus, Jascha Sohl-Dickstein, Arnaud Doucet, and Will Sussman Grathwohl. Reduce, reuse, recycle: Composi- tional generation with amery-bbssed diffusion models and meme. In International Conference on Machine Learning, pp. 8489-85 PMLR, 2023.
Bradley Efron. Tweedie's s formula and selection bias. Journal of the American Statistical Association, (66(496):1602-110 2011.
Jeffrey A Fessler and Bradley Sutton. Nonuniform fast fourier transforms using min-m min-max interpola- tion. IEEE transactions signal processing, 51 574,, 2003.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Demoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-6851, 2020.
Aapo Hyvärinen and Peter Dayan, Estimation of non normaazal statistical models by score matching. Journal of Machine Learning Research, 6(4), 2005.
Ajil Jalal, Marius Arvinte, Giannis Daras, Eric Price, Alexandros G Dimakis, and Jonathan Tamir. Robust compressed sensing mri with deep generative priors. Advances in Neural Information Processing Systems, 34, 2021.
Kyong Hwan Jin, Michael T McCann, Emmanuel Froustey, and Michael Unser. Deep convolutional neural network for inverse problems in imaging. IEEE Transactions Image Processing, 26(9): 4509-4522, 2017.
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. In Proc. NeurIPS, 2022.
Bahjat Kawar, Michael Elad, Stefano Ermon, and Jiaming Song. Denoising diffusion restoration models. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL http:///ppneeeeeeeeeeeeeereeeeee kxXvopt9pNK
Dongjun Kim, Scungjae Shin, Kyungwoo Song. Wanmo Kang, and Il-Chul Moon. Soft truncation: A universal training technique of score-based diffusion model for high precision score estimation. International conference machine learning, 2022.
Dana Knoll and David E Keyes. Jacobian-free newton krylov methods: a survey of approaches and applications. Journal of Computational Physics, 193(2)3557-377, 2004.
Anish Lahiri, Mare Klasky, Jeffrey A Fessler, and Saiprasad Ravishankar. Sparse- view cone beam et reconstruction using data-consistent supervised and adversarial learning from scarce training data. arXiv preprini arXi::2201.0933 2022.
Jörg Liesen and Zdenek Strakos. Krylov subspace methods: principles and analysis. Oxford University Press, 2013.
Guanxiong Luo, Moritz Blumenthal, Martin Heide, and Martin Uecker. Bayesian mri reconstruction with joint uncertainty estimation using diffusion models. Magnetic Resonance Medicine, 90(1): 295-311, 2023.
Michael Lustig, David Donoho, and John M Pauly. Sparse MRI: The application of compressed sensing for rapid MR imaging. Magnetic Resonance in Medicine: An Official Journal of the International Society for Magnetic Resonance in Medicine, 58(6): 182- 2007.
Neal Parikh and Stephen Boyd. Proximal algorithms. Foundations and Trends in optimization, 1(3): 127-239, 2014.
Ben Poole. Ajay Jain, Jonathan T. Barron, and Ben Mildenhall. Dreamfusion: Text-to-3d using 2d diffusion. arXiv, 2022.
11

---

Published as a conference paper at ICLR 2024
Matteo Ronchetti. Torchradon: Fast differentiable routines for computed tomography. arXiv preprint httv:2009.147880 2020.
Litu Rout, Advait Parulekar, Constantine Caramanis, and Sanjay Shakkottai. A theoretical jus- tification for image inpainting using denoising diffusion probabilistic models. arXiv preprint arXiv:2302.01217, 2023a.
Litu Rout, Negin Raoof, Giannis Daras, Constantine Caramanis, Alexandros G Dimakis, and Sanjay Shakkottai. Solving linear inverse problems provably via posterior sampling with latent diffusion models. arXiv preprint arXiv:2307.006// 2023b.
Jaaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In 9th International Conference on Learning Representations, ICLR, 2021a.
Jiaming Song, Arash Vahdat, Morteza Mardani, and Jan Kautz. Pseudoinverse-guide guudddd diffusion models for inverse problems. In International Conference on Learning Representations, 2023. URL
nwt/forummidddeeeoo _=9_gsMA8MRKK)
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution.
In Advances Neural Information Processing Systems, volume 32, 2019.
Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In 9th International Conference on Learning Representations, ICLR, 2021b.
Yang Song, Liyue Shen, Lei Xing, and Stefano Ermon. Solving inverse problems in medical imaging with score-based generative models. In International Conference on Learning Representations, 2022. URL httpps://ppeeeeva http:////...oeeeee......e/////// nett:/forummdddwaRCtyyy
Anuroop Sriram, Jure Zbontar, Tullie Murrell, Aaron Defazio, C Lawrence Zitnick, Nafissa Yakubova, Florian Knoll, and Patricia Johnson. End-to-end variational networks for accelerated MRI re- construction. In International Conference Medical Image Computing and Computer-Assisted Intervention, pp. 64-73. Springer, 2020.
Martin Uecker, Peng Lai, Mark J Murphy, Patrick Virtue, Michael Elad, John M Pauly, Shreyas S Vasanawala, and Michael Lustig. ESPIRiT-an eigenvalue approach to autocalibrating parallel MRI: where SENSE meets GRAPPA. Magnetic resonance in medicine, 9900)9901,01,,,, 2014.
Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computa- tion, 23(7):1661-160 2011.
Yinhuai Wang, Jiwen Yu, and Jian Zhang. Zero-shot image restoration using denoising diffusion null-space model. In The Eleventh International Conference Learning Representations, 2023.
URL https::/opperreiieww net//forumiidmmRiee@tttttt
Jure Zbontar, Florian Knoll, Anuroop Sriram, Tullie Murrell, Zhengnan Huang, Matthew Muckley,
Aaron Defazio, Ruben Stern, Patricia Johnson, Mary Bruno, al. fastMRI: An open dataset and benchmarks for accelerated MRI. arXiv preprint / 2018.
Kai Zhang, Wangmeng Zuo, Yunjin Chen, Deyu Meng, and Lei Zhang. Beyond a gaussian denoiser: Residual learning of deep CNN for image denoising. IEEE transactions on image processing, 26 17):3142-3155, 2017.
12

---

Published as a conference paper at ICLR 2024
A PRELIMINARIES
A.1 DIFFUSION MODELS
Let define random variable ao p(0:)) Paota( where Pdata denotes the data distri- bution. In diffusion models, we construct a continuous Gaussian perturbation kernel p(æ, mn) = Nxx;; s, of with & E [0, 1], which smooths the distribution. As t the marginal distribution Pr (xx))) is smoothed such that it approximates the Gaussian distribution, which becomes our reference distribution sample from. Using the reparametrization trick, one can directly sample
Dt S S
(O,
(29)
Diffusion models aim to revert the data noising process. Remarkably, it was shown that the data noising process and the denoising process can both be represented stochastic differential equation (SDE), governed by the score function Va, log p(æ;) (Song et al.) 2021b; Karras et al. 2022
Namely, the forward/reverse diffusion SDE can be succinctly represented as (assuming 8 for simplicity,
da = - V logp(a;;) dt = u V., log di Vojoddwe
(30)
where w, is the standard Wiener process. Here, the sign denotes the forward process, where (30) collapses to Brownian motion. With the sign, the process runs backward, and we see that the score function væ, log p(z;) governs the reverse SDE. In other words, order to run reverse diffusion sampling (i.e. generative modeling), we need access to the score function of the data distribution.
The procedure called score matching, where one tries train parametrized model Sy to approximate Vz, log p(z,) can be done through score matching (Hyvärinen & Dayan 2005). As explicit and implicit score matching methods are costly to perform, the most widely used training method the modern sense is the so-called denoising score matching (DSM) (Vincent 2011)
min Ez,. 2o.c ||
(I:) V log 1ogp(zz(10))1
(31)
which is easy to train as perturbation kernel is Gaussian. Once sg- is trained, use it as plug-in approximation the score function to plug into (30)
The score function has close relation to the posterior mean E[ which can be formally linked through Tweedie's formula
Lemma 1 (Tweedie's formula). Given
Gaussian perturbation kernel p(æ)|z0)
N(x;; stDo, I), the posterior mean is given by
1 (x, + ofV xe log (zz)()))) 8,
E[
(32)
Proof.
Va, logp(ar; ve,p(x,)
(33)
p(az)
p(x,) Va, / p(x; (x0)p(())))))) dxn
(34)
/
V Ve. dey
(35)
p(xx)
p(2,120)) Van log logp(ar, (((o)) day
(36)
p(x,)
/
p( ( Va, (x dxo
(37)
8:a0 Et dæ
p(zolz:)
(38)
sy1Ba00 [
(39)
o7
13

---

Published as a conference paper at ICLR 2024
In other words, having access to the score function is equivalent to having access to the posterior mean through time i, which we extensively leverage in this work.
A.2 KRYLOV SUBSPACE METHOODS
Consider the problem y = Az, where a. y RN,, and have square matrix A E RNN We define a fixed point iteration
Bx y, -=
y.
(40)
where D is some diagonal matrix. One can show that the iteration in (40 converges to the solution of x -'yy if and only if p(B) < 1, where p(-) denotes the spectral norm of the matrix. More specifically, we further have that
sup |0 x: 11// S p( (B).
(41)
Here, unless we know the solution x", we cannot compute the error vector
d == Fn x²
(42)
Instead, however, what we have access to is the residual
b y Ar, Ad,.
(43)
For simplicity, let D = I and thus B = I - A. The residual now becomes
b, = Ax, = Brn + y - In = In+1 Un.
(44)
Then, with 444 the Jacobi iteration reads
Øn+1 == In + b
(45)
bn+1 b, Ab, Bb..
(46)
This also implies the following
bn E Span(b, Ab,.. ,...,,,)))))
(47)
In E 3n Span (b, Ab, Ab..... A" AA)-b))
(48)
with the latter obtained by shifting the subspace of b The Krylov subspace is exactly defined by this span, iec (A) Span(b, Ab...
b).
In Krylov subspace methods, the idea is to generate a sequence of approximate solutions xx E To + K,,, (A) of Ax y. so that the sequence of the residuals b,, E (A) converges to the zero vector.
B PROOFS
(t)
Lemma (Total noise). Assuming optimality of Co the total noise w in (3) can be represented by
W. Ôt-16
(49)
for some é (O. In other words. VI- - for some é ~ N I).
23 is equivalently represented by x;-1 va-
Proof. Given the independence of the estimated é; and the stochastic e ~ ~N I) along with the
Gaussianity, the noise variance in 15 is given as 1 - 5 17 + 27 = 1 ā.-1, recovering sample from q [0)).
Proposition (Manifold Constrained Gradient). Suppose the clean data manifold M is represented as an affine subspace and assumes the uniform distribution on M. Then,
dā.
1
PM
(19)
dx
Var
8. " V a. ((t.)
PM VV, ((t))
(20)
for some > O, where PM denotes the orthogonal projection to M.
14

---

Published as a conference paper at ICLR 2024
Proof. First, provide proof for 199 vI 1 we can obtain the closed form expression of the likelihood:
Thanks to the forward sampling 21-1
sood
vda-1
(2 - (())) exp (-l1a
[
V vā(t))zll
p(a,
(50)
2(1 a())
which is Gaussian distribution. Using the Bayes' theorem, have
p(a;) = f p(a))
(51)
According to the assumption, p(æo) is uniform distribution on the subspace M. To incorporate this in (51), first compute p(æ;) by modeling p(200) as the zero-mean Gaussian distribution with the isotropic variance o² and then take the limit o + oc. More specifically, we have
p(ao) (2 exp (- |P
(52)
20²
where we use PM = 20 as 2o E M. Therefore, we have p( mo) p( (o)))
(2 a(t)))// (2x02)1/2 exp d((,,,,)))
where
d(22,
||
|PMa0ll1
+ |PMM: (a(t))Mmoll!
2(1 a(t))
2a²
2(1 a(())
IIPM300 11,11
+ c(t)]|PMass!!
2s(1)
2(1 - 2(1))(1)
and
a(t)
s(1) = 02
(53)
a(t)
V Va(t)
41
1-a(()
+ ((t) 32
(54)
G7
I-a(t)
a(t)
(55)
1 +
a(t)
I-o(t)
Therefore, after integrating out with respect to æo we have
|| x2 1 c())]|PMa const.
logp(x;))
2(1 (t)))
leading to
1 P 1- 1-a(t) PMRx
c(t)
se log)(
a(t)
Furthermore, using 55 for the uniform distribution we have lime dt) Therefore,
1
lim Va, logp(xz)) 1- a(t) Pi
000
Now, using Tweedie's 's denoising formula in 14 we have
c;
v1 ane (xx))
(56)
Var
(at (1 σ+) Vn, log (o.)(
(57)
Var
15

---

Published as a conference paper at ICLR 2024
where we use
$g. (z)) V logp(a;) = (x: -
Accordingly, we have
1 Vāc
(xt Pinaa))= = Vov
PMI.
(58)
Therefore, we have
0z.
1
PPM PM
(59)
dr.
Var
Second, we provide a proof for (00 Since ãz E M, we have 2; = PMM0r Thus, using 559) we have
dā:
- 76 ((t)
ã; - 70 dx.
V t. ((tt))
= PMI 7t PM V de (
Vā
= PM (m; - Se ((z)))
where G :- /t////// CG. Q.E.D.
For score functions trained in the context of DDPM (VP-SDE), DDIM sampling of (13) can directly be used. However, for VE-SDE that was not developed under the variational framework. is unclear how to construct sampler equivalent of DDIM for VP-SDE. As one of our goals is to enable the direct adoption of pre-trained diffusion models regardless of the framework, our goal is to devise a fast sampler tailored for VE-SDE. Now, the idea of the decomposition of DDIM steps can be easily adopted to address this issue.
First, observe that the forward conditional density for VE-SDE can be written q(æ;, (wo) M(xx (0n, I), where aN is taken to be geometric series following Song et al. (2021b). This directly leads to the following result:
Proposition 2 (VE Decomposition). The following update rule recovers sample from the marginal ((x,-1<z0) Vn E [0.
:t-1 -t w
(60)
where
à = vt + o St
(61)
W. = 0--10
82n223 + 07-1
(62)
Proof. From the equivalent parameterization as 14),
we have the following relations in VE-SDE
8: Σ, t.
(63)
of
it
(64)
Fr
c. = C2 + of3 3t.
(65)
Plugging in, we have
x:-1 t. + 01-1
329pe, + o1-1nBe.
(66)
Since the variance can now computed as
(
82212 -B22p + (0r-183) 07-1
(67)
Øt-1 g( 0r,, of I) ((x,-11/00) 1/00) by the assumption.
16

---

Published as a conference paper at ICLR 2024
Again, 22 arises from Tweedie's formula. With 17 = 1 and = original VE-SDE, whereas = 0 leads to the deterministic variation of VE-SDE, which we call VE-DDIM. We can now the usual VP-DDIM (33) VE-DDIM (60) depending the training strategy of the pre-trained score function.
-1
recover the
Here, we present main geometric observations in the VE context, which are analogous to Proposition | and Lemma 2
Their proofs are straightforward corollaries, and hence, are omitted.
Proposition 3 (VE-DDIM Decomposition). Under the same assumption of Proposition !
we have
PM
(68)
dz
V 2. ((t)
PM (ft 71 Va, ((x)))
(69)
where PMM denotes the orthogonal projection to the subspace M.
Accordingly, our DDS algorithm for VE-DDIM is as follows:
0t-1 a + w CG A, A y. Sr M), M S
(70)
a =
(71)
where CG(-) denotes the Mtstep CG for the normal equation starting from Tt.
C ALGORITHMS
In the following tables, we list all the DDS algorithms used throughout the manuscript. For simplicity, we define A, y, .) to be running M steps of conjugate gradient steps with initialization æ. For completeness, we include a pseudo-code of the CG method in Algorithm, that used throughout the work.
Algorithm 1 Conjugate Gradient (CG)
Require: A. y.y.oo Zo.
1: To + b Awo
2: Po < bg
3: for i = 0: K - do
4:
@kk Pi Apr
r
5: Db+1 t D6 + atPp
6: Tk+1 Tb aa.
r.r
7: 3k ri
8: PL+1 & b + 34P2
9: end for
10: return ZK
D DISCUSSION ON CONDITIONING
Projection type Methods that belong this class aim directly replace what have in the range space of the intermediate noisy æ; with the information from the measurement y. Two representative works that utilize projection are Chung & Ye (2022) and Song al. (022). In Chung & Ye (2022).
we use
(I A A)a + A'
(72)
where the information from y will be used to fill in the range space of A However, this is clearly problematic when considering the geometric viewpoint, as the sample may fall off the noisy manifold.
Both hard (Chung Ye 2022 soft (Song al. 2022 constraints he utilized.
17

---

Published as a conference paper at ICLR 2024
Algorithm 2 DDS (PI MRI; VP: noiseless)
Require: NV (aa))a) )
M
1: N. I))
2: for :
3: é, f- Eg- (xx)
4: Tweedie denoising
5:
I; (at, v1 / / Vāg
6:
Data consistency
7:
E CG(A" A, y, U, M))
8:
6 IO,
9:
DDIM sampling
10: Iz-1 f- vā V ā - - + nB
11: end for
20 < (æ) v1 -1- (a1))/vaar
13: return 20
Algorithm 3 DDS (PI MRI; VP: noisy)
Require: Eg., N, (or)) t= T, A, M,
xN (),
2: for = 22
3:
é f cg. (xt)
4:
Tweedie denoising
5:
i. f (x, VI- a. Va,
6:
Data consistency
7:
ACG < I -I
8:
YCG + 2z 4A"
9:
E. + CG(AcG. Bco. Le M)
10:
C~ I0,
11:
DDIM sampling
12:
I:-1 var- V ār - + 11,86
13: end for
14: 2D < (x) VT on (21))/Vo
return 30
In the Bayesian reconstruction perspective, it is natural to incorporate the gradi- ent of the likelihood as log logg(as p( (y)) V Vlog + log Hogg((yl p(y].z)). Here, while one can use
Gradient type
V log p(æx) = 8p- (æ:). V log pl y ax is intractable for all /. o 0 (Chung et al. 2023a), one has to
resort to approximations V log p xx)) In a similar spirit to Score-MRI (Chung & Ye 2022 one
can simply use gradient steps of the form
= σ ( (x))).).
(73)
where & the step size (Jalal et al. 2021 Nevertheless, V., logp(x, (y) is from Gaussian as
gets further away from 0, and hence is hard to interpret nor analyze what the gradient steps in the
direction of A (zy Axf) is leading. Albeit not in the context of MRI, a more recent approach (Chung
ct al. 2023a) proposes
(74) As 1 is the Tweedie denoised estimate and is free from Gaussian noise, 74) can thought of as minimizing the residual the noiseless data manifold M. However, must be taken since taking the gradient with respect to æ, corresponds to computing automatic differentiation through the neural net s0-, often slowing down the compute by about x2 (Chung al. 2023a
K = - #es1 - --
T2I vs. DIS For the former, the likelihood is usually given as a neural net parameterized function Po(y æ) (e.g. classifier, implicit gradient from CFG), whereas for the latter, the likelihood is given as an analytic distribution arising from some linear/non linear forward operator (Kawar et al. 2022 Chung et al. 2023a).
18

---

Published as a conference paper at ICLR 2024
Algorithm 4 DDS (3D CT recon; VP)
Require: Eg.,N, .N, [ot A, M, fA
1: ax (0,
2: NN -- torch.zeros_likee
3: WN torch. zerozzzeeos_ke(z,,
4: for = N : do
5: é < Cy. (x+)
6: Tweedie denoising
7:
&, (a, VI -,,,)/)//// Vāa
8:
Data consistency
9:
AcG + A A pD D
10:
DOG + A"y p (Z++1 Wr+1)
11:
E t AcG bcg, Lt: M)
21 t Sniip D DD + w(+1)
13:
wr + w++1 + D 2t
14:
DDIM sampling
F--1 E- v0
- + 1/3,E
16: end for
20 < -xx v1 aal (21))/ Vai
Algorithm 5 DDS (PI MRI; VE)
Require: 8g,,N, for T, A, M
1: ØN (N, o7
2: for = :
3: &, f- s-- (xx)
4:
Tweedie denoising
5:
d; xt 80- (m:))
80-
6:
Data consistency
7:
& CG(A' A, A'A'yy... y, 0r, M)
8:
1),
9:
DDIM sampling
2:-1 4- V 01-10 -0-1001 32722, 3320 + 16-16
10:
11: end for
2g < 21 + 0ss0- (a1)
13: return 20
E ALGORITHMIC DETAILS
We provide the VP counterpart of the DDS VE algorithms presented in Algorithmm16p Algorithm 2 The only differences are that the model now parameterized with eg that estimates the residual noise components and that have different noise schedule.
E. ACCELERATED MRI
As stated in section 4.1 using 2 200 NFE when using Algorithm degrades deggadee the performance due to the numerical instability. To counteract this issue, the same iteration until S N/50 k. and directly acquire the final reconstruction by Tweedie's formula:
Be 2t + o o. (IXx)
(75)
19

---

Published as a conference paper at ICLR 2024
Algorithm 6 DDS (3D CT recon; VE)
Require: sg-, .N {or NV
1=:: A, A.M, M, (
1: ax (0,
2: NN f torch. brrchzeero__lik(zzyy
3: WN torch. zerch.zzee_llie(t,)
4: for = . :
5: & <- Se. (a:)
6: Tweedie denoising
7:
&; 2, + of sg- (xt)
8:
Data consistency
9:
AcG - A pD D
10:
DoG + y'y+ p (Z++1 wt+1)
11:
E t AcG bcg, Lt. M)
21 Sniip D DD + w(+1)
13:
wr + 1++1 + D 2t
14:
DDIM sampling
15:
Iz-1 t 8 02-101 32 + 01-16
16: end for
20 < 21 + oss0 (x1)
18: return 20
(a) 3D SV-CT
(b) Multi-coi CS-MRI
forward
Reconstructed
forward measure recoustruction
Reconstructed
secoastraction .........
proj
Viewil
View
View ff2
seoppäg
Sinegram space
lespnce
Vovel space view
Meusurement
Reconstructed
Coil #1
Coil #2
Coil #N
Figure 3: Illustration of the imaging forward model used in this work. (a) 3D sparse-view CT: the forward matrix A transforms the 3ID voxel space into 2D) projections. (b) Multi-coil CS-MRI: the forward matrix A first applies Fourier transform to turn the image into k-space. Subsequently, sensitivity maps are applied element-wise product achieve multi-coil measurements. Finally, the multi-coil measurements are sub-sampled with the masks.
E.2 3D CT RECONSTRUCTION
Recall that to perform 3D reconstruction, we resort to the following optimization problem (omitting the indices for simplicity)
-- = arg min
マ AA
211 A||D
(76)
a
20

---

Published as a conference paper at ICLR 2024
One can utilize ADMM to stably solve the above problem. Here, we include the steps arrive at Algorithms 6 for completeness. Reformulating into a constrained optimization problem, we have
1
min 2.
こ ~ A 2|12/11
(77)
s.t.
== D.
(78)
Then, ADMM can be implemented as separate update steps for the primal and the dual variables
1
f++1 arg min lly
2
A
11
ar, 23 + 2ll
(79)
a;
2j+1 = arg min A||25|1 ||
+
D 2, + wl2
(80)
2
2j
j+1 w + Dzār+t 2j+1.
(81)
We have a closed-fort solution for 79
5+ (A" A pD Dz) ( y + pDD 10--00 w)),
(82)
which can be numerically solved by iterative CG
ăJ+1 CG(CGGG beg, $.j M)
(83)
AcG A A pD D=
(84)
bcg == A y + pD w)
(85)
Moreover, as (80) is in the form of proximal mapping (Parikh & Boyd 2014), we have that
Zj+1 = = S
(D. (xip(Dzx
w)
(86)
Thus, we have the following update steps
ăx+1 = AoG bos, e; M)
2j+1 = Sxip D 2 j+1 + w)
Wy+1 wi + D D. 2j+1-
Usually, such an update step is performed in an iterative fashion until convergence. However, in our algorithm for 3D reconstruction (Algorithm 6 4 we only use single iteration of ADMM per each step of denoising. This is possible as we share the primal variable z and the dual variable wo as global variables throughout the update steps. leading to proper convergence with single iteration (fast variable sharing technique (Chung et al. 2023b)).
Stabilizing Algorithm 6 As stated in section 4.2 in order to stabilize Algorithm 6 we found it beneficial to first start the DC imposing strategy with simple CG update steps without TV prior, as used in our MRI experiments. We iterate our solver starting from = N down to N/2 using standard CG updates. Then, switch to ADMM- TV scheme starting from i = N/2 down to 2.
25 Dess Core.
DBS 201
Fruudd Trum
F ADDITIONAL EXPERIMENTS
F.1 NOISE OFFSET
For this experiment, we take 50 random proton density (PD) weighted images from the fastMRI validation dataset, and add Gaussian noise oer 7.00 10 2| to the im- ages. For each noisy image, we apply the following DC step for each method
Figure 4: DDS reconstruction of CS- MRI on radial sampling trajectory. Col 1: sampling trajectory, 2: Zero filled re- construction + density compensation, 3: DDS (99 NFE), 4: ground truth.
1. Score-MRI (Chung & Ye 2022
72
21

---

Published as a conference paper at ICLR 2024
2. Jalal et al. (Jalal al. as used the implementation
2021);
73
with step size
3. DPS (Chung et al.
2023a):
74
with step size
1.0
4. DDNM (Wang et al.
2023
25
5. DDS (CG): CG applied to the Tweedie denoised
estimate, M = 5.
Once the update step is performed, the Gaussian noise level of the updated samples is estimated with the method from Chen et al. 2215). Note that as the estimation method is imperfect, there is already a gap between the ground truth noise level and the estimated noise level.
No process
DDS (CG)
Method
Score-MRI Jalal et al.
DPS DDNM
Cest
7.556
5.959
6.303
8.527
8.256
7.859
esst
0.000
1.597
1.253
0.917
0.700 0.303
Table 5: Noise offset experiment. Gaussian noise level estimated with Chen et al. (201). Real noise level: OGT = 7.00 x 10 10: Jest
7.56[x10-?]
F.2 CON-CARTESIAN MRI TRAJECTORIES
To put an emphasis on the fact that the proposed method is agnostic to the forward imaging model hand, for the first time in DIS literature, we conduct MRI reconstruction experiments on the non-Cartesian trajectory. which involves non-uniform Fast Fourier Transform (NUFFT) (Fessler & Sutton 2003). In Fig- we show that DDS is able to reconstruct high-fidelity images from radial trajectory sampling, even under aggressive accelerations.
F.3 STOCHASTICITY IN SAMPLING
For our sampling strategy. the stochasticity of the sampling process is determined by the parameter n E [o, 1]. When r) - O, the sampling becomes deterministic, whereas when v) -> the sampling becomes maximally stochastic. It is known in the literature that for unconditional sampling using low NFE, setting 7) as close to 0 leads to better performance Song et al. (202 Fig. [6] we see a similar trend when we set NFE- 20. However, when set NFE > 50, we observe that the choice of does not matter too much, as it often fluctuates within the boundary that can be as well be thought of arising from the inherent stochasticity of the sampling procedure. This is different from the observation made in (Song et al. 2021a),, which can be thought of as stemming from the conditional sampling strategy that the proposed method uses.
F.4 INSTABILITY IN VE PARAMETERIZZTTION
NFE
4000
500
200
100
50
30
Score-MRI
33.25
33.19
33.13
31.67
3,015
3.239
Chung & Ye 2022
Ours
32.07
31.16
31.99
33.69
31.79
30.40
Table 6: PSNR [db] of uniform 1D x4 acc. reconstruction with varying NFEs.
One observation that is made in this experiment, however, is that using > 200 NFEs for the proposed method degrades the performance. We find that this degradation stems from the numerical pathologies that arise when VE-SDE is combined with the parameterizing the neural network with sg. Specifically, the score function is parameterized to estimate so- (æx) ~ me me c/a,. Near t = 0. a, attains a
of
22

---

Published as a conference paper at ICLR 2024
(a) VE paramelerization s0
x xg (xx)lla
11x0 xy (xz)
-7500
17300
13500
12500
7300
7300
parameterization Es
(c) Evolution denoised estimates
(i V
5000
(i) VP
2000
1000
750
100
250
Figure 5: Evolution of the reconstruction error through time. =1.0a plot. (a) VE parameterized with 88, (b) VP parameterized with Eg. (c) Visualization of æ;
PSNR(ch)
35.0
34.0 8
Asial'
Cermaa
Supittal
Method
PSNK
SSIM
PSNK
SSIM 1
PSNR
SSEM
NFE
2023
14.92
0.935
12.48
0.94
25.82
0,832
MCE Chipa 16 Papeo
21.01
0.335
21.55
0323
21.59
0.716
31.0
Soug
Cautrer
ruse
27:00
0882
24.50
25.09
036
21.65
0773
Labil Lih 200
26.00
Zhang
26.70
0.879
25.77
0.514
23.92
0,541
=0.0
0.2
0.4
0.6
0.8
1.0
ADMM-TT
25.19
0.203
22:06
0.758
19.95
0.782
Figure Ablation study on the selection of for Algorithm 2
Table 7: Quantitative evaluation of LA-CT (90%) on the
AAPM 256x 256 test set. Bold: Best.
Asaii
Coevenl
Sezittal
Asail
Carvenl
Seaiital
Axid PSSR
Connual
Sigpital
Metked
PENT
PSNR
PSNK
PENK SSINI"
SSIN
PSNR
SSIM
8105
E/96
1.90
11.07
2.11
018
15:0
2.11
03:1
1.75
0/04
1.47
0.11
1.55
0.10
1.58
0111
3.33
0.27
1.80
0115
3.42
0.19
7.9
0.17
3.74
0.70
DPS Chargtt moteymoon
1:35
0.19
205
0.14
3.21
0.17
1.31
0.16
2.33
0.13
2.00
0.13
3.01
0.21
3.29
0.19
3.57
0.29
01155
0.18
0.13
3.69
0.30
122
1.38
0113
6118
3.08
011
1.30
0.17
1.59
0.00
1.11
6.15
7.31
310
0.34
ABMMM-T
233
0.15
0.8
2.51
0.23
288
0.38
3.44
0.22
2.00
2.70
0.16
Table 8: Standard deviation of the quantitative metrics presented in Table. which correspond to the
results for siew CT reconstruction. Mean values in Tab. 4 (Numbers in parenthesis): NFE.
very small value e.g.op = 0.01 (Kim et al. 2022). meaning that the score function has to approximate relatively large values in such regime, leading to numerical instability. This phenomenon is further illustrated in Fig. 55 (a), where the reconstruction (i.e. denoising) error has rather odd trend of jumping up and down, and completely diverging as > 0. This may be less of an issue for using samplers such PC where i; not directly used but becomes much bigger problem when
23

---

Published as a conference paper at ICLR 2024
the proposed sampler is used. In fact, for NFE > 200, we find that simply truncating the last few evolutions is necessary to yield the result reported in Tab. (See appendix E.I for details).
Such instabilities worsened when we tried scaling our experiments to complex- valued PI reconstruc- tion due to the network only being trained on magnitude images. On the other hand, the reconstruction errors for VP trained with epsilon matching have much stabler evolution of denoising reconstruc- tions, suggesting that is indeed a better fit in the context of the proposed methodology. Hence, all experiments reported hereafter use network parameterized with Eg trained within VP framework, and also by stacking real/imag part in the channel dimension to account for the complex-valued nature of the MR imagery and avoid using x2 NFE for a single level of denoising.
G EXPERIMENTAL DETAILS
G.1 DATASETS
fastMRI knee. We conduct all PI experiments with fastMRI knee dataset (Zbontar al. 2018). Following the practices from Chung & Ye (2022), among the 973 volumes of training data, we drop the first/last five slices from each volume. As the test set, we select 10 volumes from the validation dataset, which consists of 281 2D slices. While Chung et al. (2022b) used DICOM magnitude images to train the network, we used the minimum-varianne unbiased estimates (MVUE) (Jalal et al. 2021) complex-valued data by stacking the real and imaginary parts of the images into the channel dimension. When performing inference, we pre-compute the coil sensitivity maps with ESPiRIT (Uecker et al., 2014).
AAPM. AAPM 2016 CT low-dose grand challenge data leveraged in Chung et al. 2022a 2023b) is used. From the filtered backprojection (FBP) reconstruction of size 512 x 512, we resize all the images to have the size 256 x 256 in the axial dimension. We use the same volume of testing data that was used in Chung et al. (202b)). This corresponds to 448 axial slices, 256 coronal, and 256 sagittal slices in total. To generate sparse-view and limited angle measurements, we use the parallel view geometry for simplicity, with the torch-radon (Ronchetti 2020) package.
G.2 NETWORK TRAINING
VE models for both MRI/C are taken from the official github repositories (Chung & Ye 2022 Chung al., 2022a), which are both based on ncsnpp model of Score-SDE (Song al., (2021b). b). In order to train our VP epsilon matching models, we take the U-Net implementation from ADM (Dhaiiall &
Nichol 202T), and train each model for 1M iterations with the batch size of 4, initial learning rate of Te 4 on single RTX 3090 GPU. Training took about 3 days for each task.
G.3 REJECTION SAMPLING
When running Algorithm 2 in the low NFE regime (e.g. 19, 49), we see that our method sometimes yields infeasible reconstructions. Sampling 100 reconstructions for ID uniform x4 acc., see that 5% of the samples were infeasible for 19 NFE. and 3% of the samples were infeasible for 49 NFE. In such cases, we simply compute the Euclidean norm of the residual ||y A with the reconstructed sample à, and reject the sample the residual exceeds some threshold value 7. Even when we consider the additional time costs that arise from re-sampling the rejected reconstructions, we still achieve dramatic acceleration to previous methods Song et al. (2022); Chung & Ye (2022).
G.4 COMPARISON METHODS
G.4.1 ACCELERATED MRI
Score-MRI (Chung& & Ye, 2022) We use the official pre-trained mode with 2000 PC sampling. Note that for Pl, this amounts to running the sampling procedure per coil. When reducing the number of NFE presented in Fig. I we use linear discretization with wider bins.
24

---

Published as a conference paper at ICLR 2024
Jalal et al. 22021). As the official pre-trained model is trained fastMRI brain MVUE images, in order to periorm fair comparison, we train the NCSN v2 model with the same fastMRI knee MVUE images for 1M iterations in the setting identical to when training the model for the proposed method. For inference, we follow the default annealing step sizes as proposed the original paper. We use Langevin dynamics steps per noise scale for 700 discretizations, which amounts total of 2100 NFE. When reducing the number of NFE presented in Fig. I we keep the Langevin steps intact, and use linear discretization with wider bins.
Varnet (Sriram et al., 2020), U-Net. We train the supervised learning-based methods with Gaussian ID subsampling as performed in Chung & Ye (2022), adhering to the official implementation and the default settings of the original work.
TV. We use the implementation sigpy Eegpy.mri.og maps with ESPiRiT (Uecker et al. 2014). The parameters for the optimizer are found via grid search on 50 validation images.
www.mrrpppppp.ptt.mmmm app. TotalVariationf with calibrated sensitivity
G.4.2 3D CT RECONSTRUCTION
DiffusionMBIR (Chung al. 2023b), MCG (Chung al. 2022a). Both methods use the same score function as provided in the official repository For both DiffusionMBIR and Score-CT, CT, we use 2000 PC sampler (4000 NFE). For DiffusionMBIR. we set 10.0. p 0.04, which is the advised parameter setting for the AAPM 256x256 dataset. For Chung et al., we use the iterating ART projections as used in the comparison study Chung et al. (2023b).
S
Lahiri et al. (2022). We use two stages of 3D U-Net based CNN architectures. For each greedy optimization process, we train the network for 300 epochs. For CG, we use 30 iterations at each stage. Networks were trained with the Adam optimizer with a static learning rate of le 4, batch size of 2.
FBPConvNet Jin et al. (017), Zhang et al. (007)). Both methods utilize the same network architecture and the same training strategy, only differing in the application (SV, LA). We use standard 2D U-Net architecture and train the models with 300 epochs. Networks were trained with the Adam optimizer with learning rate of le 4, batch size 8.
ADMM-TV. Following the protocol of Chung et al. (2023b), we optimize the following objective
x
arg
arg min デ
1| y11 x|| 2,11
(87)
with D : Dz, Dy, D.). and is solved with standard ADMM and CG. Hyper-parameters are set identical Chung et al. (2023b).
H QUALITATIVE RESULTS
"https://ittub.c "hets://gthub.cm//mikggauu www.h//www...ommmmomonnngggggyyy
http:/hhhhhhooooooooaaa http://gghtut..oo////aaaa/CCCCCC diffussii
25

---

Published as a conference paper at ICLR 2024
VarNet
Score-MRi
24.91 8.71 715
2B. 6B D 715 (8)
30, 28 0 723
44 76tl
0.715
11.3910.783
21. GB.G 0.892
20.23 0.60
36.32 0.BnS
37. 07.161
Figure 7: Comparison of parallel imaging reconstruction results. (a) subsampling mask (1st row: UniformIDx4. 4, 2nd row: MouusianlDx$s 3rd row: G 4th row; variable density poisson disc x8), (b) U-Net (Zbontar et al. 2018). (c) E2E- VarNet (Sriram et al. (2020), (d) Score-
MRI (Chung & Ye 2022 (4000 2 e NFE), (e) DDS (49 NFE), ground truth.
26

---

Published as a conference paper at ICLR 2024
(50)
(1000)
(49)
04
22.31 D
220 D.BS
25.58
20. G.BAA
32.35/0.855 0.854
24.1910.704
Figure 8: Comparison of noisy (a p.0a)la imaging reconstruction results. (a) subsampling mask (1st row: uniform x 4, 2nd row: Gaussian Dxx4, 3rd row: Gaussian2D x8, 4th row: variable density poisson disc 8). (b) Zero-filled, (c) DPS (Chung et al. 2023a) (50 NFE), (d) DPS (Chung jet al. 2023a) (1000 NFE), (e) DDS (49 NFE), (f) ground truth. Numbers the top right corners denote PSNR and SSIM, respectively.
27

---

Published as a conference paper at ICLR 2024
[q]
13. / 4044
(b)
20.65 (0.688
(c)
27.88 .8 0.860
(d)
33.34/0.938 (e)
34.67 009553
(
3D sparse CT (Eview)
14.19/00 19.0.887
27.7670.912
28. 04
28.988
85, 0.
11.41
27.00 0.832
0.932
(e)
0.932
3D limiter: angle CT (90s)
13.78
18.73/0.826
21 21.25+0.010
23.15
86 0.054
0.099
Figure 9: Comparison of 3D CT reconstruction results. (Top): 8-view SV-CT, (Bottom): 90" LA-CT. (a) FBP, (b) Lahiri al., (c) Chung et al., (d) DiffusionMBIIR (4000 NFE), (e) DDS (49 NFE), (f) ground truth. Numbers in top right corners denote PSNR and SSIM, respectively.
28


Reviews:
[{"review_id": "0xszRuO3tZ", "reviewer": "Reviewer_Q4Mc", "summary": "This work proposes the Decomposed Diffusion Sampling (DDS) method as a Diffusion model-based Inverse problem Solvers (DIS) for inverse problems in medical imaging. The work is based on the observation that a diffusion posterior sampling (DPS) with the manifold constrained gradient (MCG)  is equivalent to one-step projected gradient on the tangent space at the “denoised\" data by Tweedie’s formula, this work\nprovides multi-step update scheme on the tangent space using Krylov subspace methods. The experiments show that performing numerical optimization schemes on the denoised representation is superior to the previous methods of imposing DC. Further, the work devises a fast sampler based on DDIM that works well for both VE/VP settings. With extensive experiments on multi-coil MRI reconstruction and 3D CT reconstruction, it was shown that DDS achieves superior quality while being ≥ ×80 faster than the previous DIS.", "questions": "1. Do we need the  assumption that the clean data manifold is an affine subspace ? How about general curved manifold ?\n2. Why Krylov space method guarantee to stay in the tangent space ?\n3. Please elaborate the geometric view of diffusion model more. Is the diffusion process on the manifold or in the ambient Euclidean space ?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The main strength of the work is the novel theoretic insight: the geometric interpretation that the diffusion posterior sampling (DPS) with the manifold constrained gradient (MCG)  is one-step projection to the tangent space of the clean data manifold, and using Conjugate Gradient type method, one can achieve multiple steps projection within the tangent space. All the theoretic proofs are given in details, all the lemmas are well formulated and clearly explained. The experimental results are convincing.", "weaknesses": "The work involves both Krylov space theory and diffusion model, it will be more helpful for general audience to give brief overview of both theories in the appendix. Especially. the key observation : DPS with MCG is a one-step projection to the tangent space of the clean data manifold. Furthermore, it needs more explanation for the problem : why Krylov space method can guarantee to stay in the tangent space.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "l931VDUqBN", "reviewer": "Reviewer_wm9K", "summary": "In this paper, the authors aim to solve the inverse problem through diffusion model and propose the conjugate gradient (CG) method based on its connection to Krylov subspace. Specifically, in each time step, they update the noisy estimation based on the CG update and finally obtain the inverse results. Experiments show that the proposed method works well.", "questions": "Please see above", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Using CG method in each time step is reasonable and can mitigate the difficulty of solving the original problem.\n- The results seem promissing.", "weaknesses": "- How to initialize $x_0$, from Gaussian noise?\n- In equation (19), what is $A_t$ and $y_t$, are they the same with $A$ and $y$? If so, will this cause some bias or problem?\n- For different steps, will $l$ in equation (22) be fixed? If so, how to get an optimal $l$?\n- In table 1, what's the meaning of mask pattern?\n- In table 2, it seems that with more steps, the performance drops. Any explanation?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "A9t50XQ0EK", "reviewer": "Reviewer_tJRQ", "summary": "This paper proposes a novel and efficient diffusion sampling strategy that synergistically combine the diffusion sampling and Krylov subspace methods.\nSpecifically, it prove that if the tangent space at a denoised sample by Tweedie’s formula forms a Krylov subspace, then the CG initialized with the denoised data ensures the data consistency update to remain in the tangent space.\nThis eliminates the necessity of computing the manifold-constrained gradient (MCG), leading to a more efficient diffusion sampling method. \nThe proposed method achieves state-of-the-art  reconstruction quality on challenging real-world medical inverse imaging problems, and more than 80 times faster inference time than the  previous state-of-the-art method.", "questions": "1.In page 3, \" the m-th order Krylov subspace is defined by equation (2)\", this \"m-th\" is not shown in equation (2).\n2.For high-dimensional data, the computational complexity of the Conjugate Gradient (CG) algorithm is significant, and each step of inverse sampling requires solving using the CG algorithm. \nDoes this affect the model's efficiency?\n3.Which has a higher computational complexity, this method or the classical accelerated diffusion models(e.g., DDIM)?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "A novel and efficient diffusion sampling strategy is proposed, synergistically combining diffusion sampling and Krylov subspace methods. \nIt certifies that if the tangent space at a denoised sample by Tweedie’s formula forms a Krylov subspace, \nthen the Conjugate Gradient initialized with the denoised data ensures that the data consistency update remains in the tangent space. \nThis also eliminates the need to compute the manifold-constrained gradient.", "weaknesses": "1.The detailed process of M-step CG update is lacking.\n2.In page 7, I think clarifying Equation (29) will help the reader understand the problem setting.\n3.An important point of this paper is the negation of the necessity to compute the manifold-constrained gradient, \nbut there is a lack of elaboration on using Tweedie's formula to obtain the tangent space for gradient calculation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "nyCuzeeHBm", "reviewer": "Reviewer_YQ7m", "summary": "In this paper, the authors propose to integrate classical Krylov subspace methods, widely used in large linear inverse problems, with state-of-the-art (SOTA) diffusion models applied to medical image reconstruction. An efficient diffusion sampling strategy is proposed that, if the tangent space at a denoised sample forms a Krylov subspace, allows the Conjugate Gradient (CG) method to maintain data consistency updates within this subspace. This eliminates the need for the computationally expensive manifold-constrained gradient (MCG) reducing computation time. The approach achieves  SOTA reconstruction quality in large scale inverse imaging problems, such as multi-coil MRI and 3D CT reconstruction.", "questions": "1.\tIn section 2, it is stated that in the forward model (1): y=Ax, we can assume A square without losing generality because, if not, we can consider the modified equation (2): A*y=A*Ax, because its solution is also a solution of problem (1). However, equation (1) has infinite number of solutions when A has fewer rows than columns as in the considered in the image reconstruction scenarios, and the solution of eq. (2) is unique if A is full rank matrix. It is not clear why the unique solution of eq. (2) is the most preferable solution over all possible solutions of eq. (1). Could you please clarify about this?\n\n2.\tThe proposed method has a limitation since it only considers linear forward models with the form y=Ax and uses Krylov method to minimize l(x) = ||y - Ax||^2. However, just after eq. (6) it is mentioned that Krylov method can be extended to non-linear problems and the example l(x) = 0.5||y - Ax||^2 is given, which is confusing since it is the same optimization problem as in the linear model y=Ax. Can you clarify what do you mean by extending Krylov method to nonlinear problems? Can thus you apply your method to an inverse problem having a nonlinear forward model?\n\n3.\tThe methods requires that the tangent space on the clean manifold at a denoised sample forms a Krylov subspace and experimental results suggests that such assumption is met. What is the intuition behind this assumption? Can you provide some conditions for such assumption become true?\n\n4.\tRegarding the reported PSNR and SSIM values: Are these average values? If so, how many samples have you consider computing the average? Can you also report the variability of these measures to derive conclusions about the comparison of methods? Are the differences between methods statistically significative in some sense?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "•\tThe idea of applying Krylov subspace methods for efficient sampling strategies in stable diffusion models is novel and significant since practical applications such as medical image reconstructions are usually large inverse problems.\n\n•\tThe organization of the paper is clear with useful introductions to the relevant topics such as diffusion models and Krylov subspace methods.\n\n•\tComparison with various SOTA methods is provided showing that the proposed method achieves higher quality results with significant fewer number of iterations.", "weaknesses": "Main weaknesses:\n\n•\tThe comparison provided primarily centers on state-of-the-art (SOTA) diffusion models for reconstructing images from undersampled data and only a comparison with a the naïve Total Variation (TV) is included. However, I believe it lacks a comparison with the classical Compressed Sensing (CS) approach, which relies on sparse image representations and has shown success, especially in MRI applications (as in the classical reference [*]). It's unclear whether SOTA diffusion models can compete with traditional CS algorithms. Perhaps referencing recent studies that compare diffusion models with classical CS methods would offer a more comprehensive perspective on the field's current status.\n\n•\tIt is not clear why the Krylov method, originally developed for square matrix A can be successfully applied to the image reconstruction problem when A has fewer rows than columns. See my question number 1 below.\n\n•\tThe proposed method has a limitation since it only considers linear forward models with the form y=Ax (see my question number 2 below).\n\n•\tThe methods requires that the tangent space on the clean manifold at a denoised sample forms a Krylov subspace and experimental results suggests that such assumption is met. However, an intuition about why the tangent space forms a Krylov subspace is missing. It would be useful to have identified conditions for such assumption become true (see my question number 3 below).\n\n•\tExperiments: Quality results are expressed in terms of PSNR and SSIM values, but it is not clear if these are averaged values and the number of instances considered in the corresponding average. The experiments should also report variability of the results, for example showing standard deviation for a more comprehensive comparison of different methods (See my question number 4).\n\nMinor comments:\n\n-\tFig. 2 caption should indicate what are the numbers at the bottom of each subimage, which I think they are PSNR/SSIM\n\n-\tAcronym NFE is not defined in the paper, please introduce its definition (Number of Forward Evaluations?)\n\n-\tJust before eq. (3): “two subspace K and L” -> “two subspaces K and L”\n\n-\tJust after eq. (7): “for given the optimization problem” -> “for a given the optimization problem”\n\n-\tPage 6: “guranteed” -> “guaranteed”\n\n-\tProblem setting considers only the noiseless case (equation (29)), however experiments in Fig. 2 shows some results in the noisy case. Please include the noisy case in the problem setting.\n\n-\tIn the problem setting, s^{(1)}, s^{(2)}, …, s^{(c)} are not defined. For single-coil case s^{(1)}=I but it is not clear how they are defined in the multi-coil case. Please, clarify.\n\n-\tPC sampler is not defined.\n\n[*] Lustig, M., Donoho, D. and Pauly, J.M. (2007), Sparse MRI: The application of compressed sensing for rapid MR imaging. Magn. Reson. Med., 58: 1182-1195. https://doi.org/10.1002/mrm.21391", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)