PaperlayerView exact revision on arXiv

Source-linked AI summary

Decomposed Diffusion Sampler for Accelerating Large-Scale Inverse Problems

Hyungjin Chung, Suhyeon Lee, Jong Chul Ye

arXiv:2303.05754v3cs.LGcs.AIcs.CVstat.ML

TL;DR

Diffusion inverse-problem solvers offer adaptable priors but remain slow, and it is unclear how to combine them effectively with fast Krylov optimization. The paper proposes DDS, which performs CG updates on the tangent space of Tweedie-denoised samples and uses DDIM sampling. DDS achieves strong reconstruction quality on multi-coil MRI and 3D CT while being at least 80 times faster than prior diffusion inverse-problem solvers.

  • Problem

    Diffusion models generalize across forward models but have slow inference, while existing accelerated consistency strategies are not successful for complex large-scale medical imaging problems.

  • Method

    DDS combines DDIM diffusion sampling with CG updates that remain on the tangent space of Tweedie-denoised data, eliminating the need for computationally demanding MCG.

  • Results

    DDS achieves superior reconstruction quality on multi-coil MRI and 3D CT while being ≥×80 faster than previous diffusion inverse-problem solvers.

  • Takeaways & Limitations

    DDS provides a fast diffusion sampling strategy for challenging large-scale medical inverse problems in both VE and VP settings.

  • Takeaways & Limitations

    The authors call for rigorous validation and clinical collaboration before applying the method in real-world medical settings.

Abstract

from arXiv · show

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

1 INTRODUCTION

Diffusion models provide adaptable generative priors for inverse problems, but their inference cost remains a major obstacle. DDS combines diffusion sampling with Krylov methods to accelerate data-consistent reconstruction while retaining applicability across medical imaging tasks.

  • Motivation: Diffusion models can combine a parameterized prior with a measurement likelihood, allowing one pretrained network to support different forward models.This posterior formulation is useful for inverse problems whose acquisition physics vary across applications.
  • Motivation: Medical imaging reconstruction requires adaptable algorithms because acquisition schemes vary across vendors, sequences, and measurement patterns.The passages contrast this requirement with supervised methods that can overfit to training measurement types.
  • Motivation: Slow inference time is a critical drawback of diffusion models, and existing DDIM-based consistency strategies can fail for complex large-scale medical imaging operators.The cited examples include parallel-imaging compressed-sensing MRI and 3D modalities.
  • Classical optimization: Krylov methods seek approximate solutions in affine subspaces generated by successive applications of the linear operator to the residual, with CG minimizing residuals in that subspace.Their rapid convergence makes them relevant to large-scale inverse problems.
  • Proposed approach: DDS uses Krylov updates on the tangent space at a Tweedie-denoised sample, avoiding computationally demanding manifold-constrained gradients while supporting VP and VE sampling.The method performs multiple economical CG steps during each ancestral diffusion sampling stage.
  • Reported outcome: 20∼50 NFE and ×80 ∼200 acceleration are reported for DDS on multi-coil MRI and 3D CT reconstruction.Representative comparisons are identified in Figures 1 and 2.

2 BACKGROUND

The background introduces projection-based Krylov solvers for linear systems and diffusion models as denoising generative processes. It also describes DDIM as a faster alternative to ancestral sampling and motivates combining multi-step optimization with diffusion updates.

  • Krylov subspace methods: Projection-based Krylov methods seek an approximate solution in a chosen subspace while enforcing an orthogonality condition on the residual.For a nonzero estimate, the update can be expressed as finding a correction δ in the chosen subspace.
  • Krylov subspace methods: CG uses the same Krylov subspace for its search and constraint spaces, with K = L = K_l = Span(b, Ab, · · · , A^l−1b).The residual vector is b = y − A x̂.
  • Krylov subspace methods: Krylov methods can address nonlinear inverse problems by linearizing near the current solution and applying standard Krylov solvers to the resulting problem.The background presents Newton-Krylov methods as one such extension.
  • Combining the methods: The paper uses the distinction between gradient iterations and Krylov subspace expansion to motivate computationally efficient multi-step optimization within diffusion sampling.This connection supports replacing a single projected-gradient-style update with multiple Krylov steps.
  • Diffusion models: Diffusion models construct hierarchical latent variables through a forward noising process and train denoisers across noise levels for reverse generation.The background covers VP and VE formulations, including a Brownian-motion form for VE sampling.
  • Diffusion sampling: Ancestral diffusion sampling iteratively generates less noisy variables, while DDIM accelerates sampling by matching marginal distributions q(x_t|x_0).DDIM introduces a stochasticity parameter η, with η = 0 yielding deterministic sampling and η = 1 corresponding to the fully stochastic endpoint under the stated setting.
  • Diffusion sampling: The DDIM noise component matches the forward marginal, and its difference from DDPM lies in the degree of dependence on the deterministic noise estimate.Intermediate stochasticity values are allowed between the deterministic and stochastic extremes.

3 DECOMPOSED DIFFUSION SAMPLING

Conditional diffusion alternates reverse diffusion with data-consistency optimization, but manifold-constrained gradients are expensive and unstable. DDS replaces them with CG updates that remain in the tangent space under a Krylov-subspace condition, while extending to noisy restoration and 3D CT.

  • Conditional diffusion: Conditional diffusion alternates reverse diffusion with operations minimizing the data-consistency loss while seeking the clean manifold.For linear inverse problems, the loss is ℓ(x) = ∥y − Ax∥2/2.
  • Limitations of MCG: MCG computation requires expensive backpropagation and is often unstable.
  • Manifold-constrained gradients: Under an affine clean-manifold assumption, the DPS update corresponds to projected gradient descent on the clean manifold.The projection interpretation follows from Proposition 1 and its orthogonal-projection operator PM.
  • Krylov-subspace formulation: A tangent space represented by a Krylov subspace lets M-step CG updates remain within that tangent space when M ≤ l.The construction uses Kt,l = Span(b, Ab, · · · , A^(l−1)b), with b = y − A x̂t.
  • DDS algorithm: DDS therefore avoids MCG computation by using standard CG initialized at the denoised sample.The method compares CG-based updates with DDNM’s pseudoinverse update, which does not guarantee tangent-space membership.
  • Empirical rationale: CG outperforms naive projections for large-scale inverse problems and avoids cumbersome step-size tuning.The paper also reports that DDS empirically stays close to the noisy manifold even when its Krylov assumptions cannot be guaranteed.
  • Extensions: DDS extends to noisy restoration through proximal CG and to 3D CT by optimizing conditioning in the clean manifold.For noisy restoration, CG minimizes a proximal objective weighted by γ; for 3D CT, the additional finite-difference prior acts along the axial direction.

4 EXPERIMENTS

Experiments evaluate DDS on accelerated and noisy multi-coil MRI and 3D CT reconstruction, comparing sampling strategies, parameterizations, and baselines. DDS improves reconstruction quality while substantially reducing inference time across these settings.

  • Experimental setup: Experiments cover accelerated multi-coil MRI and sparse-view or limited-angle 3D CT reconstruction under varied measurement settings.MRI tests use diverse sub-sampling patterns; CT tests use sparse-view and limited-angle reconstruction on the AAPM 256×256 dataset.
  • Accelerated MRI: DDS outperforms DDNM by a few dB in PSNR, while five CG iterations per denoising step add only about 0.2 seconds at 50 NFE.The reported comparison is 4.51 seconds analytically versus 4.71 seconds with CG(5).
  • Accelerated MRI: DDS exceeds the 4000-NFE Score-MRI reconstruction quality with 100 NFE under VE sampling and still produces decent reconstructions at 30 NFE.The VE setting becomes numerically unstable with large NFE, as discussed in the appendix.
  • Accelerated MRI: DDS generalizes across MRI sub-sampling patterns, including non-cartesian NUFFT measurements, while supervised baselines fail dramatically on 2D patterns.The experiments use VP parameterization and report reconstruction across four sub-sampling types.
  • Accelerated MRI: DDS sets the state of the art in most MRI cases below 100 NFE, outperforming Score-MRI at 49 NFE and matching it at 19 NFE.The reported acceleration remains ×80–×200 even before accounting for Score-MRI’s additional multi-coil NFE factor.
  • Noisy MRI: For noisy multi-coil MRI with σ = 0.05 complex Gaussian noise, DDS far outperforms DPS at 1000 NFE while being about ×40 faster.DDS handles the noisy problem with CG without requiring singular value decomposition.
  • 3D CT reconstruction: For sparse-view CT, DDS reaches quality comparable to DiffusionMBIR with 100 versus 4000 NFE and achieves state-of-the-art results at ≤49 NFE with VP parameterization.Reported wall-clock times are approximately 25 minutes at 49 NFE and 10 minutes at 19 NFE, compared with two days for DiffusionMBIR.

5 CONCLUSION

DDS combines diffusion sampling with conjugate-gradient optimization on denoised representations, achieving superior medical inverse-imaging quality while being at least 80 times faster than previous diffusion inverse solvers.

  • 5 CONCLUSION: DDS combines a geometric diffusion view with conjugate-gradient solvers and a DDIM-based sampler for VE and VP settings.The method performs numerical optimization on denoised representations and provides tailored algorithms across settings.
  • 5 CONCLUSION: Experiments on multi-coil MRI and 3D CT reconstruction report superior quality and inference at least 80 times faster than the previous diffusion inverse solver.
  • 5 CONCLUSION: The experiments use publicly available, fully anonymized medical data, while the authors call for rigorous validation and clinical collaboration before real-world application.The stated concerns include patient-data privacy and possible misinterpretation of generated images.
  • 5 CONCLUSION: Score matching trains a parameterized model to approximate the data-distribution score, which can then be used in the reverse diffusion SDE.

A.2 KRYLOV SUBSPACE METHODS

Krylov subspace methods construct approximate solutions in an affine subspace generated by repeated applications of a matrix, using residuals to obtain rapid convergence for large inverse problems.

  • A.2 KRYLOV SUBSPACE METHODS: Because the unknown solution is unavailable, the method uses the residual rather than the error vector to guide iterations.
  • A.2 KRYLOV SUBSPACE METHODS: The Krylov subspace K_n(A) is the span of b, Ab, through A^n b.
  • A.2 KRYLOV SUBSPACE METHODS: Krylov methods seek approximate solutions x_n in x_0 + K_n(A) while residuals b_n in K_{n+1}(A) converge toward zero.

B PROOFS

The proofs develop DDS from diffusion-manifold geometry: Tweedie denoising supplies a clean estimate, while decomposed VE/VP updates and CG enforce data consistency through tangent-space operations.

  • B PROOFS: Under an affine clean-data manifold with a uniform distribution, the manifold-constrained gradient is expressed using orthogonal projections onto the manifold and its complement.
  • B PROOFS: The VE decomposition provides an update rule that recovers a sample from the marginal q(x_{t−1}|x_0) for every η in [0, 1].
  • B PROOFS: Setting η = 0 yields the deterministic VE-DDIM variation, while η = 1 recovers the original VE-SDE.
  • B PROOFS: DDS applies CG initialized at the Tweedie-denoised estimate, with the CG step defined as an M-step solve of the normal equation.
  • B PROOFS: Projection-based conditioning can fill the range space of A† but may move samples off the noisy manifold.
  • B PROOFS: Gradient-based conditioning can require backpropagation through the score network, slowing computation by about ×2, whereas DDS uses CG on the denoised representation.
  • B PROOFS: The VP DDS algorithms use an epsilon parameterization and a different noise schedule from the VE algorithms.

E.1 ACCELERATED MRI

The accelerated reconstruction pipeline combines physics-specific forward models with CG, ADMM, and shared variables to stabilize and speed medical inverse imaging; high NFE can destabilize the VE MRI procedure.

  • E.1 ACCELERATED MRI: At least 200 NFE degrades Algorithm 5, so the method repeats iterations only through k = N/50 before obtaining the final reconstruction with Tweedie’s formula.
  • E.1 ACCELERATED MRI: For multi-coil CS-MRI, A applies a Fourier transform, multiplies by coil sensitivity maps, and subsamples measurements with masks.
  • E.1 ACCELERATED MRI: For 3D sparse-view CT, A maps the 3D voxel space into 2D projections.
  • E.1 ACCELERATED MRI: ADMM reformulates the 3D reconstruction into constrained primal and dual updates, with the primal subproblem numerically solved by iterative CG.
  • E.1 ACCELERATED MRI: The 3D algorithm uses one ADMM iteration per denoising step because global sharing of primal and dual variables supports convergence.
  • E.1 ACCELERATED MRI: For stability, Algorithm 6 first uses standard CG updates and then switches to an ADMM-TV scheme during later denoising steps.

F ADDITIONAL EXPERIMENTS

Additional experiments examine DDS under noisy MRI conditions, non-Cartesian sampling, and different stochasticity settings. The results support high-fidelity radial-trajectory reconstruction and characterize how η affects sampling behavior.

  • Noise and implementation: 50 noisy proton-density images from fastMRI were evaluated with Gaussian noise at σGT = 7.00[×10−2].The experiment applied each method’s data-consistency step to the noisy images.
  • Noise and implementation: DDS implementations include CG applied to the Tweedie denoised estimate with M = 5, alongside DPS and other comparison procedures.The reported noise-level estimate is σnp est = 7.56[×10−2] versus real noise σGT = 7.00[×10−2].
  • Non-Cartesian MRI: DDS reconstructs high-fidelity CS-MRI images from radial non-Cartesian sampling trajectories, even under aggressive acceleration.Figure 4 compares the sampling trajectory, density-compensated zero-filled reconstruction, DDS at 99 NFE, and ground truth.
  • Stochasticity ablation: When NFE = 20, lower η improves performance, whereas for NFE ≥50, the choice of η has little effect.η controls stochasticity: η approaching 0 is deterministic, while η approaching 1 is maximally stochastic.

F.4 INSTABILITY IN VE PARAMETERIZATION

The VE parameterization becomes unstable at high NFE, with reconstruction error diverging near the terminal noise regime. The reported experiments therefore favor VP parameterization for the proposed sampler.

  • VE instability: Using ≥200 NFEs degrades the proposed method’s performance because VE-SDE with sθ develops numerical pathologies.The associated reconstruction error can jump, oscillate, and diverge as t →0.
  • VE instability: For NFE > 200, truncating the final evolutions is necessary to obtain the reported Table 6 result.This provides a practical workaround for the VE instability at high sampling budgets.
  • Parameterization choice: VP models trained with epsilon matching show stabler denoising-error evolution, so all subsequent experiments use VP with ϵθ.The instability worsens for complex-valued parallel-imaging reconstruction because the network was trained only on magnitude images.

G.1 DATASETS

The study evaluates DDS on fastMRI knee MRI and AAPM 2016 CT data, using multiple reconstruction baselines and specified training and optimization protocols. Low-NFE MRI sampling can occasionally produce infeasible reconstructions that are rejected and resampled.

  • Datasets: MRI experiments use fastMRI knee data, with 10 validation volumes comprising 281 test slices after removing five slices from each volume end.The training data include 973 volumes, and the reconstruction network uses complex-valued MVUE images.
  • Datasets: CT experiments use one AAPM 2016 test volume resized to 256 × 256, including sparse-view and limited-angle geometries.The test volume contains 448 axial, 256 coronal, and 256 sagittal slices.
  • Rejection sampling: At low NFE, 5% of samples were infeasible at 19 NFE and 3% at 49 NFE, so samples exceeding a residual threshold were rejected.The residual test uses the Euclidean norm of y − Aˆx, including resampling costs in the acceleration comparison.
  • Baselines and protocols: DDS is compared with diffusion, supervised, compressed-sensing, and optimization baselines using the listed implementations and training protocols.The baselines include Score-MRI, DPS, U-Net, E2E-VarNet, DiffusionMBIR, FBPConvNet, ADMM-TV, CG, and related methods.

H QUALITATIVE RESULTS

Qualitative results compare DDS with established MRI and CT reconstruction methods across sampling masks, noise levels, and CT acquisition geometries. The figures report reconstruction outputs alongside ground truth and, where specified, PSNR and SSIM.

  • Parallel MRI: Parallel-imaging comparisons span four subsampling masks and contrast U-Net, E2E-VarNet, Score-MRI, DDS at 49 NFE, and ground truth.The masks include uniform 1D ×4, Gaussian 1D ×8, Gaussian 2D ×8, and variable-density Poisson-disc ×8.
  • Noisy parallel MRI: Noisy parallel-imaging comparisons use σ = 0.05 and include zero-filled, DPS at 50 or 1000 NFE, DDS at 49 NFE, and ground truth.PSNR and SSIM values are shown in the image corners.
  • 3D CT: 3D CT comparisons cover 8-view sparse-view CT and 90° limited-angle CT against FBP, learned baselines, DiffusionMBIR at 4000 NFE, DDS at 49 NFE, and ground truth.PSNR and SSIM values are displayed in the top-right corners.
Loading 2303.05754v3…