PaperlayerView exact revision on arXiv

Source-linked AI summary

Solving 3D Inverse Problems using Pre-trained 2D Diffusion Models

Hyungjin Chung, Dohoon Ryu, Michael T. McCann, Marc L. Klasky, Jong Chul Ye

arXiv:2211.10655v1cs.CVcs.AIcs.LG

TL;DR

3D inverse problems remain difficult for diffusion models because full-dimensional generation is computationally expensive and 3D priors are data hungry. DiffusionMBIR combines pre-trained 2D diffusion sampling with z-direction TV regularization and MBIR updates, achieving state-of-the-art reconstruction across three medical imaging tasks, including two-view CT.

  • Problem

    Diffusion models had not been extended effectively to 3D inverse problems because full-dimensional generation demands excessive memory and computation, while 3D priors require substantial data.

  • Method

    DiffusionMBIR runs 2D diffusion denoising independently across slices and augments it with z-direction TV regularization, slice aggregation, and efficient ADMM updates.

  • Results

    DiffusionMBIR achieves state-of-the-art reconstructions on sparse-view CT, limited-angle CT, and compressed-sensing MRI, including accurate sparse-view CT with as few as two views.

  • Takeaways & Limitations

    The method produces coherent 3D reconstructions across axes using 2D diffusion models, while supporting high-dimensional voxel reconstruction with limited training data.

Abstract

from arXiv · show

Diffusion models have emerged as the new state-of-the-art generative model with high quality samples, with intriguing properties such as mode coverage and high flexibility. They have also been shown to be effective inverse problem solvers, acting as the prior of the distribution, while the information of the forward model can be granted at the sampling stage. Nonetheless, as the generative process remains in the same high dimensional (i.e. identical to data dimension) space, the models have not been extended to 3D inverse problems due to the extremely high memory and computational cost. In this paper, we combine the ideas from the conventional model-based iterative reconstruction with the modern diffusion models, which leads to a highly effective method for solving 3D medical image reconstruction tasks such as sparse-view tomography, limited angle tomography, compressed sensing MRI from pre-trained 2D diffusion models. In essence, we propose to augment the 2D diffusion prior with a model-based prior in the remaining direction at test time, such that one can achieve coherent reconstructions across all dimensions. Our method can be run in a single commodity GPU, and establishes the new state-of-the-art, showing that the proposed method can perform reconstructions of high fidelity and accuracy even in the most extreme cases (e.g. 2-view 3D tomography). We further reveal that the generalization capacity of the proposed method is surprisingly high, and can be used to reconstruct volumes that are entirely different from the training dataset.

1. Introduction

Diffusion models provide powerful data-driven priors for inverse problems, but their full-dimensional generative process makes 3D reconstruction computationally prohibitive. DiffusionMBIR combines pre-trained 2D diffusion models with model-based reconstruction to produce coherent 3D results efficiently.

  • Diffusion priors: Diffusion models learn data distributions through score functions and can generate high-quality, mode-covering samples with robust training.Their training reduces to a denoising mean squared error objective.
  • Motivation: Pre-trained diffusion models have effectively solved several inverse problems, including sparse-view CT and compressed sensing MRI, but prior work focused on 2D imaging.The 2D focus reflects the high-dimensional generative constraint involved in iterative denoising.
  • Motivation: 3D diffusion is difficult because generation remains in the full data dimension, creating prohibitive memory and computational costs and requiring data-hungry 3D priors.Training a 3D model would typically require thousands of volumes, unlike 2D models trained with substantially less data.
  • Model-based reconstruction: Model-based iterative reconstruction formulates imaging as weighted least squares with data consistency and regularization, while total variation preserves edges and imposes smoothness.The paper notes that TV is too simplistic to fully model image appearance compared with data-driven priors.
  • Proposed method: DiffusionMBIR inserts MBIR optimization into diffusion sampling, augmenting a 2D data-driven prior with a z-direction TV prior and aggregating slices for data consistency.The method uses ADMM updates and variable sharing to reduce optimization cost while supporting SV-CT, LA-CT, and CS-MRI.
  • Proposed method: The proposed method operates on voxel representations, targets dimensions above 256^3, and can be trained with fewer than ten 3D volumes.These contributions emphasize scalability and reduced data requirements for 3D reconstruction.

2. Background

The paper reviews model-based and diffusion-based formulations for inverse reconstruction, then explains why extending diffusion models directly to 3D is computationally difficult. It motivates combining diffusion priors with model-based reconstruction for volumetric imaging.

  • Model-based iterative reconstruction: 3D inverse reconstruction can be formulated with a linear measurement model and a regularized least-squares objective.The measurement y may be a sinogram or k-space, while A represents a Radon or Fourier transform and R supplies regularization.
  • Model-based iterative reconstruction: TV regularization preserves edges and imposes smoothness but is considered too simplistic to fully model image appearance.The paper contrasts this conventional prior with data-driven priors from modern machine learning.
  • Score-based diffusion models: Diffusion models generate data-space samples by reversing a noising process, using a learned score function during numerical integration.The score can be trained with denoising score matching, which acts as residual denoising across multiple noise scales.
  • 3D diffusion: Direct 3D diffusion is difficult because reverse diffusion remains in the full voxel dimension, with volumes beyond 64^3 requiring excessive GPU memory.Point-cloud alternatives are efficient but too sparse for tomographic reconstruction, which requires accurate interior estimation.
  • 3D diffusion: A latent 3D diffusion approach models 160 × 224 × 160 volumes through 20 × 28 × 20 latents but requires 1000 synthetic 3D training volumes.The passage also notes that applying latent diffusion to inverse problems is not straightforward.
  • Diffusion inverse problems: Posterior sampling can alternate a diffusion-prior update with projection onto the measurement-consistent subspace.The method is described using a numerical reverse-SDE solver and a projection operator, with stochasticity implicitly defined.

3. DiffusionMBIR

DiffusionMBIR combines slice-wise 2D diffusion denoising with a 3D ADMM update that imposes z-directional regularization and data consistency. Variable sharing and sub-batching address convergence and memory constraints.

  • 3.1. Main idea: Slice-wise 2D diffusion alone can produce incoherent reconstructions because it ignores inter-dependency between slices.The problem becomes especially pronounced for sparser view angles.
  • 3.1. Main idea: DiffusionMBIR replaces 2D projection with a 3D ADMM update while retaining slice-by-slice diffusion denoising.This combines the diffusion prior in each slice with volumetric model-based optimization.
  • 3.1. Main idea: The regularizer applies the ℓ1 norm only to finite differences along z because the neural network already models the xy-plane prior.The resulting prior augmentation targets spatial correlation in the remaining direction and produces coherent 3D reconstructions across all axes.
  • 3.2. Algorithmic steps: The algorithm alternates SDE integration, ADMM updates, and inner conjugate-gradient iterations.The ADMM step augments z-directional TV and imposes consistency after parallel slice denoising.
  • 3.2. Algorithmic steps: Reinitializing ADMM variables at every SDE iteration slows convergence and discards information between adjacent diffusion states.The proposed variable-sharing strategy keeps the updated primal-associated variables throughout the SDE iterations.
  • 3.2. Algorithmic steps: Sub-batching enables denoising when a full-volume neural-network pass exceeds commodity-GPU memory.The sub-batches are aggregated before applying the full-volume ADMM step, yielding the same solution as the corresponding algorithm.

4. Experimental setup

The experiments evaluate DiffusionMBIR on sparse-view CT, limited-angle CT, and compressed-sensing MRI, using held-out or out-of-distribution data and comparisons with reconstruction baselines.

  • CT experiments use AAPM 2016 low-dose challenge volumes, with all but one volume used to train the 2D score function and one held out for testing.
  • The study evaluates sparse-view CT, limited-angle CT, and compressed-sensing MRI using PSNR and SSIM averaged over the evaluation data.
  • Sparse-view CT: Figure 3 compares 8-view SV-CT reconstructions across axial, sagittal, and coronal slices for FBP, ADMM-TV, prior methods, DiffusionMBIR, and ground truth.
  • Comparison methods: For CT comparisons, the study includes Chung et al.’s diffusion approach, implemented with POCS/ART because its MCG step requires 10GB of VRAM for one 256×256 slice.
  • Compressed-sensing MRI: The CS-MRI experiments compare against score-MRI, DuDoRNet, and U-Net, training on fastMRI knee data and testing on BRATS Flair data.

5. Results

DiffusionMBIR improves reconstruction quality across sparse-view CT, limited-angle CT, and compressed-sensing MRI, while also showing strong out-of-distribution behavior and a useful directional TV-prior design.

  • Sparse-view CT: DiffusionMBIR outperforms baselines by large margins in most sparse-view CT settings and preserves structural connectivity across axial, sagittal, and coronal views.
  • Limited-angle CT: In 90° limited-angle CT, DiffusionMBIR improves over diffusion and fully supervised methods while robustly preserving structure across missing-view reconstructions.
  • Compressed-sensing MRI: At approximately 2× acceleration in compressed-sensing MRI, the proposed method shows large improvements over prior methods.
  • Out-of-distribution performance: With 8-view SV-CT, the method produces high-fidelity reconstructions for out-of-distribution ellipsis and randomly placed-sphere phantoms despite their different anatomy.
  • Choice of augmented prior: The method imposes the TV prior only in the redundant z-direction, leaving the xy-plane to the diffusion prior.
  • Choice of augmented prior: An ablation finds that TV(xyz) produces blurrier images than TV(z), suggesting that applying TV in the xy-plane shifts results away from the trained diffusion prior.

6. Conclusion

DiffusionMBIR performs 3D medical image reconstruction using a 2D diffusion model augmented with a TV prior and integrated with ADMM iterations. It achieves state-of-the-art results across several tasks, including accurate two-view sparse-view CT and reconstruction of out-of-distribution data.

  • DiffusionMBIR combines a 2D diffusion model with a TV prior operating along the redundant z direction for 3D reconstruction.
  • DiffusionMBIR integrates diffusion sampling steps with ADMM iterations for efficient 3D medical image reconstruction.
  • The method achieves state-of-the-art reconstructions on sparse-view CT, limited-angle CT, and compressed-sensing MRI.
  • Two-view sparse-view CT reconstructions remain accurate under the paper's reported extreme setting.
  • DiffusionMBIR reconstructs out-of-distribution volumes that differ substantially from the training data.

Supplementary Material

The supplementary material formulates the ADMM-TV framework for TV-regularized weighted least squares and summarizes its alternating optimization steps. It also identifies the AAPM test volume used in the dataset specification.

  • A. ADMM-TV: The ADMM-TV framework solves a TV-regularized weighted least-squares reconstruction problem.
  • A. ADMM-TV: The optimization splits variables and uses the scaled formulation of ADMM to alternate between updates.
  • A. ADMM-TV: The smooth subproblem has a closed-form solution, with conjugate gradients available instead of directly computing a matrix inverse.
  • A. ADMM-TV: The proximal mapping of the ℓ1 norm is implemented using the soft-thresholding operator.
  • A. ADMM-TV: The scaled dual-variable update is w+ = w + Dzx+ − z+, and the algorithmic details are provided in Algorithm 2.
  • B.1. Dataset: The AAPM dataset specification designates the L097 volume for testing and the remaining volumes for training.

B.1. Dataset

The experiments use AAPM CT, BRATS brain MRI, and fastMRI knee data, with the supplied materials describing reconstruction setup and qualitative comparisons for four-view and two-view sparse-view CT.

  • AAPM: The AAPM data come from the 2016 CT low-dose grand challenge and use fan-beam geometry with varying acquisition parameters.
  • AAPM: AAPM reconstruction uses single-slice rebinning from helical cone-beam projections followed by standard filtered backprojection.
  • Sparse-view CT: Figure 6 compares four-view sparse-view CT reconstructions across FBP, ADMM-TV, two prior methods, the proposed method, and ground truth.
  • Sparse-view CT: Figure 7 compares two-view sparse-view CT reconstructions from FBP and the proposed method against ground truth.
  • BRATS: The BRATS test volume is the first FLAIR volume with matrix size 240 × 240 × 154, while methods were trained using the separate fastMRI 2019 knee database.

B.2. Details of network training

The CT score network is trained with an unmodified NCSN++ architecture using Adam on a single RTX 3090 GPU. Training uses a linear warm-up schedule and lasts about one and a half weeks.

  • The CT score function uses an unmodified NCSN++ network trained with λ = σ2(t) and ε = 10^-5.
  • Adam training uses β1 = 0.9, β2 = 0.999, a linear warm-up reaching 2 × 10^-4 at step 5000, and batch size 2.
  • Training runs for 200 epochs on a single RTX 3090 GPU and takes about one and a half weeks.

B.3. Comparison methods

The paper compares DiffusionMBIR with reconstruction methods using diffusion, CNN, and TV-based iterative optimization components.

  • Chung et al. combines a pre-trained diffusion score function with ART applied after every PC update step.
  • Lahiri et al. uses two-stage 3D CNNs trained with slabs, implemented with a standard U-Net architecture.
  • FBPConvNet uses the same U-Net architecture as Lahiri et al., but trains it only on 2D images.
  • ADMM-TV minimizes an isotropic total-variation objective with 30 outer ADMM iterations and 20 inner conjugate-gradient iterations.The parameters are selected through coarse and visual grid searches, with different values reported for SV-CT and LA-CT.

C.1. Additional experimental results

Additional experiments examine extreme sparse-view tomography and reconstruction performance as the number of measured views changes. The method remains effective at very low view counts but degrades sharply below four views.

  • The method reconstructs 4-view and even 2-view sparse-view tomography, producing measurement-feasible and realistic multiple reconstructions.The 2-view regime contains too little measurement information for accurate reconstruction without a strong generative prior.
  • Number of views vs. performance: Performance remains close to its higher-view level down to 8 views, with only a modest drop in reconstruction quality.
  • Number of views vs. performance: Performance caps above 16 measured views and heavily degrades below 4 views.
  • CS-MRI: The CS-MRI experiment reports axial, sagittal, and coronal reconstructions for several baselines, the proposed method, and ground truth.PSNR and SSIM values are displayed for the compared reconstructions, alongside the sub-sampling mask.
  • Augmented-prior ablation: The augmented-prior ablation compares isotropic TV in xyz, the proposed TV prior in z, and ground truth.
Loading 2211.10655v1…