DiffusionBlend: Learning 3D Image Prior through Position-aware Diffusion Score Blending for 3D Computed Tomography Reconstruction

Q: Explain why wrong structures exist in the reconstruction

A:

• With ultra-sparse views (such as 4 views), CT reconstruction is extremely difficult since very few measurements are taken, so many traditional reconstruction algorithms completely fail in this scenario as demonstrated in Figure 3. Also, non-deep learning methods generate lots of structures that do not exist even with as much as 20 views as demonstrated in Figure 2 in the rebuttal pdf. To account for this missing information with ultra-sparse views, DiffusionBlend(++) generates many image structures that represent the ground truth image well, and some other structures that may look different from ground truth, due to the extremely big challenge of this ultra-sparse view reconstruction task. Note that our method is still significantly better than all compared baselines in such a challenging task.
• With increasing number of views, more measurements are taken, and it is more likely we can reconstruct accurate structures (as mentioned by reviewer Vnw7). We provide reconstruction examples in the rebuttal pdf; results show that the image structures look almost identical to the ground truth with 40 views or more.

Q: The method's purpose is to decrease the cost of memory, time. Show the time consuming and memory consuming

A:

• Firstly, the main purpose of this work is, for the first time, to investigate how to learn a 3D patch diffusion prior that can improve ultra-sparse 3D CT reconstruction performance. The abstract mentions the difficulty of training a diffusion model on 3D full volumetric data (e.g. simply cannot fit into a 48G A40 GPU), which motivates us to train a 3D patch diffusion model which is much more computational efficient than a full 3D model, while achieving impressive results.
• We have provided the inference time of our method and other baselines in Table 9 of our paper on the 500 slices of the LDCT test set. Additionally, we also provide the training time, training memory, inference memory in the rebuttal pdf. Overall our method is efficient since it only uses a 3-channel diffusion model. Surprisingly, our method has better inference time than SBTV, while do not require much more memory than SBTV. The reason is that unlike SBTV which requires expensive 3D TV computation, our method does not require any external regularizations, so the inference speed is faster.

Q: Whether more conditional slides will make numerical performance better?

A: We perform experiments that train our conditional model (DiffusionBlend) on 0, 2, 4, 6 conditional slices and evaluate on the LIDC dataset with 8 projections. Results show that including conditional slices increases the performance significantly, but adding more conditional slices is not guaranteed to further improve the performance and the improvement may be marginal. To explain why the improvement is marginal with more conditioning slices, we can think about the iterative process of diffusion reverse sampling [1,2]. For example, if using 2 conditional slices, for one specific slice, the first iteration conditions on 2 neighboring slices, but its neighbors also conditions on 2 additional slices. At the second iteration, that slice still conditions on its neighbors (which already had information from 2 additional slices), so it can get the information from the condition of 4 slices, by tracing back to the first iteration. In this way, we can show that with sufficient number of iterations, one specific slice can take the information (condition) on every other slice. Since we have sufficient number of iterations, either 2,4 or 6 conditional slices enables us to condition one slice on all other slices eventually.

This method enables us to learn the 3D prior very efficiently with minimal number of conditioning slices.

Q: The adjacent slices of different 3D patches still cannot be updated simultaneously

A: Previous work [1,3] demonstrate that each step of summation, gradient descent, or even ADMM updates can be split into each step of diffusion reverse sampling process while achieving satisfying performance. Here even though the adjacency slices may not be updated simultaneously at consecutive reverse sampling iterations, it is similar to a Monte Carlo average of the score of the distribution of different partitions as demonstrated in Eq.7. So the goal is not to compute the score of the 3D patch distribution as in Eq.5, but to approximate the score of the ground truth distribution p(x) by averaging the scores of the distributions of multiple partitions.

[1] Chung, Hyungjin, et al. "Solving 3d inverse problems using pre-trained 2d diffusion models." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023

[2] Chung, Hyungjin, Suhyeon Lee, and Jong Chul Ye. "Decomposed Diffusion Sampler for Accelerating Large-Scale Inverse Problems." The Twelfth International Conference on Learning Representations.

[3] Lee, Suhyeon, et al. "Improving 3D imaging with pre-trained perpendicular 2D diffusion models." Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023.

[4] Song, Yang, et al. "Solving Inverse Problems in Medical Imaging with Score-Based Generative Models." International Conference on Learning Representations.

[5] Chung, Hyungjin, et al. "Improving diffusion models for inverse problems using manifold constraints." Advances in Neural Information Processing Systems 35 (2022): 25683-25696.

[6] Ye, Siqi, et al. "SPULTRA: Low-dose CT image reconstruction with joint statistical and learned image models." IEEE transactions on medical imaging 39.3 (2019): 729-741.

[7] Goldstein, Tom, and Stanley Osher. "The split Bregman method for L1-regularized problems." SIAM journal on imaging sciences 2.2 (2009): 323-343.

Thank you for providing the valuable feedback. Below, we address the concern: First, we address concerns regarding to additional experiments:

Q: Evaluate method's performance at various angles.

A:

• We provide results on [20, 40, 60, 80, 100] views for our method (DiffusionBlend++) as well as baselines with those angles (including the SB(Split-Bregman)TV, CGLS, and SIRT baselines) for the LDCT dataset in the rebuttal pdf as well in the table below. Results show that our method outperforms the baselines significantly for every angle we evaluated on. From Figure 1,2,3 in the rebuttal pdf DiffusionBlend++ starts to reconstruct images very close to the ground truth with 20 projections or more, but other baselines such as SIRT and CGLS still struggle to get a satisfying reconstruction with 60 views or more. The following table shows the reconstruction performance on the coronal plane of different methods. We observe that DiffusionBlend++ (ours) has a significant margin above baselines for every view. Note that DiffusionBlend++ still outperforms DDS2D significantly with >40 views, which demonstrates that our 3D prior is still very useful even with much more views.
• The motivation of this paper is to exploit the 3D data-driven prior (instead of external 3D regularization or learning another 2D prior on another plane) for medical image reconstruction. The hypothesis is that with the aid of the 3D prior, we can have better reconstruction quality comparing to only using 2D priors. Many previous works leverage 2D diffusion priors for sparse-view reconstructions ( > 10 views), but only a few works exploit reconstructions with fewer than 10 views. By this consideration, we do not provide results for more than 10 projection in the main paper. Nevertheless, we will add additional results with more projection angles and more traditional baselines in the camera-ready version.
  Table.1 Coronal plane results on LDCT test set | Method |4 views| 6 views| 8 views | 20 views | 40 views | 60 views | 80 views | 100 views | |------------------|----------|----------|----------|----------|----------|----------|-----------|-----------| | FBP | 11.16/0.236 | 13.18/0.268 | 14.64/0.325 | 20.22/0.594 | 26.09/0.792 | 29.97/0.886 | 32.89/0.933 | 35.16/0.959 | | SIRT | 22.12/0.823 | 23.54/0.841 | 24.50/0.851 | 27.72/0.882 | 30.82/0.918 | 33.34/0.945 | 35.56/0.962 | 37.66/0.974 | | CGLS | 22.11/0.823 | 23.49/0.841 | 24.44/0.850 | 27.67/0.882 | 30.66/0.918 | 32.96/0.944 | 34.82/0.960 | 36.36/0.971 | | SBTV | 23.33/0.849 | 25.83/0.881 | 27.85/0.896 | 31.51/0.938 | 34.83/0.966 | 36.74/0.963 | 37.59/0.973 | 38.97/0.979 | | DDS2D | 30.25/0.916 | 32.33/0.939 | 33.64/0.950 | 35.34/0.961 | 38.26/0.977 | 40.07/0.984 | 41.19/0.987 | 41.88/0.989 | | DDS | 30.89/0.932 | 32.95/0.879 | 33.97/0.935 | 35.81/0.967 | 37.59/0.975 | 38.26/0.978 | 39.07/0.980 | 39.54/0.983 | | DiffusionBlend++ | 34.27/0.955 | 36.66/0.963 | 37.87/0.968 | 40.56/0.980 | 42.53/0.987 | 43.66/0.990 | 44.42/0.992 | 44.95/0.993 |

Q: Comparison with classic CT reconstruction methods, not just FBP

A: Table.1 and Figure.1,2,3 in the rebuttal pdf show more results on additional classic baselines. Note that DiffusionBlend++ does not need to tune any hyperparameter when the number of projections changes, but this is a requirement for some traditional methods.

• SBTV: We implement this algorithm with variables splitting of 3D anisotropic TV regularization (Dz, Dx, and Dy). We first check number of iterations, note that the performance converges with around 30 iterations. We did a grid search of hyperparameters on 9 validation images (not in the test set) for every projection angles. For example: the searching for 20 projections is conducted as below:

• SIRT: This algorithm iteratively updates the reconstruction based on the residual between projection of the reconstruction and the GT. It only has the number of iterations as its hyper-parameter. We note that during inference, PSNR increases with more iterations, but saturates later. So we set the total number iterations to be 1000, with an early stopping threshold of 1e-6 between two consecutive iterations.

• CGLS: This algorithm uses conjugate gradient for solving least square problems. In our case, we use $CG(A^TA + \rho x^Tx, A^Ty)$, $\rho$ is set to be 1e-4 based on grid search for numerical stability. We tune the number of iterations on validation set, and find that performance saturates at around 25 iterations.

Q: Whether the projections have noise?

A: Following approaches in [1,2,3,4,5], we do not add noise in our projections since the noise is very small for normal dose. Nevertheless, we also run additional experiments with 8-view reconstruction for LDCT test set in the low-dose setting with Poisson-Gaussian noise. The pre-log noise model is given by $y_i = \text{Poisson}(I_0 \exp(-Ax_i)) + N(0, \sigma^2)$. Following [6] to account for same dose level, we set $I_0 = 10^6$, $\sigma = 5$. Results as below show that our method is robust to noise.

We thank the reviewer for providing the encouraging feedback! Below, we address the concern:

Q: Discuss related work with other regularization methods

A: We are aware of other works that apply external regularization between 2D slices with diffusion models such as [1,2] as we cited in our paper. Some other works train a xy plane prior and a yz plane prior [3] (also cited) for medical image reconstruction. Nevertheless, to the best of our knowledge, our work is the first one that uses only one 3D generative prior for medical image reconstruction without any additional regularization or priors. We will add more discussions on related works in our camera-ready version.

Q: Release code

A: Thank you for the encouraging feedback! We are in the process of refactoring our code, and will release it ASAP after acceptance.

Q: Difference between the score-med paper

A: The problems we solve and our motivations are different from the score-med paper [4]. The score-med paper is one of the first works that leverages a 2D diffusion prior for solving 2D medical image reconstruction problems. Our work is the first work that investigates whether using a 3D patch diffusion prior can lead to improvement over a 2D diffusion prior for 3D medical image reconstruction. The answer to the question is positive as demonstrated in our paper. Also our baseline "DDS 2D" is a very similar approach to the score-med paper. We show that our method outperforms DDS 2D by a significant margin.

Q: Window center and window width

A: We follow [1,2,3,4]'s approach for processing CT images, to see bone structures better. Our window center is 100 HU, and window length is 2800 HU ([-1300 HU, 1500 HU]). In the rebuttal pdf, we also provide the visualization of our reconstruction in a much narrower window [-150 HU, 250 HU], and we observe that our method can reconstruct the fine details and structures of images very accurately with more projections.

Q: However, diffusion models are still slow compared to other methods both traditional and deep learning and non-diffusion way.

A: We demonstrate that our method and other diffusion-based approaches can be faster than TV-based iterative traditional approaches (SBTV (one of the best non-deep-learning approach)) in our rebuttal pdf. The reason is the computation of 3DTV can be costly, but our approach does not need any external regularization so it can be faster.

[1] Chung, Hyungjin, et al. "Solving 3d inverse problems using pre-trained 2d diffusion models." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023

[2] Chung, Hyungjin, Suhyeon Lee, and Jong Chul Ye. "Decomposed Diffusion Sampler for Accelerating Large-Scale Inverse Problems." The Twelfth International Conference on Learning Representations.

[3] Lee, Suhyeon, et al. "Improving 3D imaging with pre-trained perpendicular 2D diffusion models." Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023.

[4] Song, Yang, et al. "Solving Inverse Problems in Medical Imaging with Score-Based Generative Models." International Conference on Learning Representations.

Dear Reviewer FN1q,

We really appreciate your time for writing the review and providing the encouraging feedback. We want to highlight some novelties in our algorithm design in additional to our motivations and problem settings as the author-reviewer discussion window is closing.

The key novelty in our paper mainly focuses on designing an efficient 3D generative prior by learning from the 3D volume for CT reconstruction. We will revise our paper to highlight the our key novelty, and clearly state our contribution. In our paper, we have the several key novel designs:

1. stack each 2D slice and each adjacency slice of a 3D medical image as a multi-channel image, and use a diffusion prior to learn the distribution of the patch of the adjacency slices. To our best knowledge, we are the first to propose this learning strategy in medical image reconstruction.
2. Propose a random blending algorithm for randomly partitioning and blending adjacency patches, which demonstrates to significantly improve slice inter-smoothness. To our best knowledge, we are the first to propose this method for medical image reconstruction.
3. More importantly, we propose a novel jumping-slice patch idea, which enable our diffusion prior to learn the long-range dependency instead of focusing on adjacency slices. We treat non-adjacency slices as a wide patch and learn the distribution of 3D patches with different thickness. This method enables us to work on larger volumes. To our best knowledge, we are the first to propose this method. We demonstrate that we can directly learn the long-range dependency via jumping slices. Results show this design significantly reduces image artifacts as demonstrated in Figure.4 in our main paper.

The novelty of the score-med paper [1] lies in its inverse problem solving technique and posterior sampling, which is not the main focus of our paper. We are more focusing on learning the 3D prior distribution, and we propose several novel designs to achieve this goal. Thanks for your consideration again, hope this comment resolves some of your concerns. Feel free to let us know if there is any remaining question about the manuscript and we will try our best to answer.

[1] Song, Yang, et al. "Solving Inverse Problems in Medical Imaging with Score-Based Generative Models." International Conference on Learning Representations.

Thank you for providing the valuable feedback. Below, we address the concern:

Q: Extending the method to a broader field

A: Our method is flexible and does not assume any specific data modality and forward models, so it should work for other modalities such as natural 3D images or videos, and for other inverse problems as well. Nevertheless, the results on other modalities and tasks are beyond the scope of this work, but we tested unconditional generation of 3D CT images which show good generation quality and inter-slice smoothness. We will include this experiment in the camera-ready version as reference. Thus, we think that the inter-slice smoothness of the learned prior from our method could be potential to work for videos as well to obtain a inter-frame coherence, which can be a very interesting future works. We will add this discussion in the camera-ready version as well.