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

Diffusion models face significant challenges when employed for large-scale medical image reconstruction in real practice such as 3D Computed Tomography (CT). Due to the demanding memory, time, and data requirements, it is difficult to train a diffusion model directly on the entire volume of high-dimensional data to obtain an efficient 3D diffusion prior. Existing works utilizing diffusion priors on single 2D image slice with hand-crafted cross-slice regularization would sacrifice the z-axis consistency, which results in severe artifacts along the z-axis. In this work, we propose a novel framework that enables learning the 3D image prior through position-aware 3D-patch diffusion score blending for reconstructing large-scale 3D medical images. To the best of our knowledge, we are the first to utilize a 3D-patch diffusion prior for 3D medical image reconstruction. Extensive experiments on sparse view and limited angle CT reconstruction show that our DiffusionBlend method significantly outperforms previous methods and achieves state-of-the-art performance on real-world CT reconstruction problems with high-dimensional 3D image (i.e., $256 \times 256 \times 500$). Our algorithm also comes with better or comparable computational efficiency than previous state-of-the-art methods.

Paper

References (55)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer FN1q7/10 · confidence 5/52024-07-09

Summary

This paper used the diffusion model to solve the widely known inverse problem in computed tomography. The paper is well-written and well-organized. I reckon this paper has two main contributions: 1. The first diffusion model in CT considered z-axis consistency and used 3d data as neural network input. 2. A SOTA results with astonishing visualization and metrics results.

Strengths

1. The first diffusion model in CT considered z-axis consistency and used 3d data as neural network input. 2. A SOTA results with astonishing visualization and metrics results.

Weaknesses

1. I admit this is the first paper I read using the diffusion model while considering z-axis consistency. However, the y-axis has been widely discussed and solved using various types of regularization between 2D slices or just using 3D data before. I consider the author just used a 3D way here to solve this point. Yeah, it is new for the diffusion model, but I do think it is less attractive to me considering the works already made before. 2. The results achieved by the authors, particularly in sparse-view reconstruction using only four angles, are indeed remarkable. The quality of the reconstruction with such limited data is surprising and commendable. Given the significance of these results, it would be beneficial for the community if the authors consider sharing their code upon acceptance to enable further research and validation. 3. The paper's novelty is somewhat overshadowed by the work "Solving Inverse Problems in Medical Imaging with Score-Based Generative Models," which the authors have cited. It would be beneficial for the authors to clearly state the innovations and contributions of their work, especially in light of the foundational problem settings of inverse problem solving for medical imaging already addressed by the cited paper.

Questions

1. What is your window center and width for CT images in the main body of the paper? I think the current window width value you selected for the main part of the paper is quite wide, considering using a narrower one. I think some details are not so clear in this setting, though I can tell your result is better than others.

Rating

7

Confidence

5

Soundness

3

Presentation

3

Contribution

2

Limitations

1. Although the author used some way to speed up the process both for training and sampling. However, diffusion models are still slow compared to other methods both traditional and deep learning and non-diffusion way.

Authorsrebuttal2024-08-14

Looking for your feedback and input

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.

Reviewer FN1q2024-08-14

I think the author's reply solved my doubts, I will improve my score. I suggest that the git repo be added to the final version, regardless of whether the paper is accepted.

Reviewer Vnw76/10 · confidence 4/52024-07-09

Summary

1. The method proposed a method that learns the 3D-patch image prior incorporating the cross-slice dependency. 2. The method achieves state-of-the-art reconstruction results for 3D volumetric imaging for the task of ultra-sparse-view and limited-angle 3D CT reconstruction on "AAPM 2016 CT challenge" dataset and "LIDC-IDRI" dataset.

Strengths

1. The probabilistic modeling of the 3D volume taking into account neighboring slices is novel. 2. The paper describes the author's motivation and implementation plan very well.

Weaknesses

1. Sparse angle CT imaging is typically defined as using fewer than 100 angles. The paper should evaluate the method's performance at various angles (e.g., 20, 40, 60, 80, and 100 angles). Demonstrating that this method outperforms the comparison method in PSNR/SSIM at fewer than 10 angles does not necessarily indicate better performance at other sparse angles (e.g., 20, 40, 60, 80, and 100 angles). 2. The comparison should include classic CT reconstruction methods, not just FBP, as FBP lacks prior information and performs poorly among traditional methods. Suggested comparison methods include: 2.1 SIRT (Simultaneous Iterative Reconstruction Technique) algorithm 2.2 Conjugate Gradient Least Squares (CGLS) algorithm 2.3 Split-Bregman (SB) Total Variation: Goldstein, T. and Osher, S., 2009. The split Bregman method for L1-regularized problems. SIAM journal on imaging sciences, 2(2) Because the traditional iterative method also involves the adjustment of hyperparameters, for a fair comparison, please first use grid search to search the hyperparameters of the traditional method to obtain the best performance of the traditional method, and then compare the performance of the traditional method with the performance of this paper. Also, please report the extent of the grid search. Only in this way can it can be fully evaluated in the numerous algorithms for CT image reconstruction. 3. Figure 3 shows a large number of structures that do not exist in CT images, which will interfere with the doctor's diagnosis. Please explain the reasons for this phenomenon and how to avoid it. (For example, increase the number of angles, which is why I am very concerned about increasing the angles). If increasing the angles can avoid the artifacts that come with the generative model on a large scale, the potential for medical applications will increase. 4. As the abstract said, the method's purpose is to decrease the cost of memory, time. In order to prove that this method has achieved the original intention of it, please show the time consuming and memory consuming during training phase and testing phase. And compare the memory and time consuming between different methods.

Questions

1. If the "conditioned slices" in equation (2) increase("conditioned slices" is: x[:,:,i-j:i-1] and x[:,:,i+1:i+j]), will the numerical performance better? If it is better, how does the performance gain change as the number of adjacent slices increases? 2. The paper mentions using a different partition so that the previous border slices can be included in another partition. However, if the method creates a new partition and then uses this new partition to create 3D patches, the adjacent slices of different 3D patches still cannot be updated simultaneously when the algorithm attempts to update the slices based on the scores calculated in each 3D patch. So, my question is, if the method uses the joint distribution modelling method (Eq. (5)), how can it avoid the situation I mentioned? 3. For the experimental part, The reconstruction mentioned here is the reconstruction from the simulated projection to the 3D CT volume or the reconstruction from the real CT projection(real measurement) to the 3D CT volume? If it is a simulated projection, is simulated noise added to the simulated projection? (Note: The noise I am referring to here is not the noise added in the diffusion model to train the neural network, but the noise added to simulate the real projection, which is generally a combination of Poisson noise and Gaussian noise.)

Rating

6

Confidence

4

Soundness

1

Presentation

3

Contribution

3

Limitations

The artifacts generated by the diffusion model may lead to misjudgment in the doctor's diagnosis. Please provide methods to avoid artifact, such as increasing the angle or the amount of training data?

Authorsrebuttal2024-08-07

Rebuttal (Part 2)

> **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. | Method | Training memory | Training time | Inference memory | Inference time | |------| ------| ------| ------| ------| |SBTV [7] | - | -| 6045MB | 62 mins | |DiffusionMBIR [1] | 29043MB | 47 hours | 22062MB | 23 hours | |DDS [2] | 27040MB | 12 hours | 20608MB | 48 minutes | |DiffuisonBlend++ (Ours) | 35384MB | 4.5 hours | 9976MB | 32 minutes |

Authorsrebuttal2024-08-07

Rebuttal (Part 3)

> **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. | Conditioning slices | PSNR/SSIM (8 views)| |------| ------| |0 | 30.98/0.894| |2 | 33.73/0.933| |4 | 33.32/0.932 | |6 | 33.53/0.936 | > **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.

Reviewer YiqS5/10 · confidence 4/52024-07-11

Summary

This paper proposes a novel method for learning 3D diffusion priors for CT reconstruction, which does not require large-scale data or computational resources. It presents two approaches: DiffusionBlend and DiffusionBlend++. The former learns a specific frame given adjacent slices, while the latter learns a 3D patch. Experimental results demonstrate that both methods are efficient and outperform previous works.

Strengths

1. This method enables 3D reconstruction without the need for large-scale data and resources. 2. It achieves excellent inter-slice consistency without relying on external regularizations. 3. Additionally, this method outperforms existing baselines.

Weaknesses

The only concern is that this method may be too simplistic to be extended to a broader field. For example, I believe that inter-slice smoothness cannot be guaranteed if this method is applied to videos. However, given that this method is proposed for CT reconstruction, I don't think this issue warrants the rejection of the paper.

Questions

I don't have questions about this paper.

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors have discussed the limitations.

Area Chair Jwyf2024-08-08

Paper discussion and rating finalization

Dear Reviewers, Can you please have a look at the reports of the other reviewers and also the rebuttal from the authors and respond to their questions, if available, then discuss any issues of your concern, finalise and reach a consensus about the rating of the paper before the deadline of the next Tuesday, 13th August? Thank you very much for your time, effort and contribution to the organization of NeurIPS 2024, AC

Reviewer Vnw72024-08-10

I think the author's reply solved some of my doubts, I will improve my score. I suggest that the supplementary experiments be added to the final version, regardless of whether the paper is accepted or not.

Authorsrebuttal2024-08-10

Thanks for the review and reading our rebuttal. Your review is crucial for us to improve our manuscript. We will definitely add all the supplementary experiments discussed in the rebuttal period to our paper. Feel free to let us know if there is any remaining question about the manuscript and we will try our best to answer.

Authorsrebuttal2024-08-10

Thanks for catching this typo on the checklist. We confirm that all the datasets used in this work are from public sources with no new data introduced. We apologize for misunderstanding this checklist question, and will change the answer to [NA] with the justification: "there is no new dataset introduced in the paper, and all data are from public sources."

Area Chair Jwyf2024-08-11

Reminder of the deadline

Dear Reviewers, While the deadline, Tuesday, 13th August, is approaching, can you please check at your early convenience the rebuttal from the authors, make or require further clarifications, if necessary, and interact with the authors over any further issues/concerns you may still have, and finalise the rating of the paper soon. Even though you have no further issues/concerns, you may want to acknowledge the responses to your questions in the rebuttal from the authors. Thank you very much for your time, effort and contribution to the organization of NeurIPS 2024, AC

Authorsrebuttal2024-08-13

Further clarifications on more general applications of our method

Dear Reviewer YiqS, We present some additional preliminary analysis on whether our method is still feasible on video data, and the results look promising. Hope that this new evidence can be helpful to answer your question. To further address your concern about a broader application of the proposed method, we first compute the inter-slice consistency for unconditional generation of 3D CT generation with / without random blending. In addition, we also perform more experiments on unconditional video generation to see whether our algorithm can maintain inter-slice (frame) consistency for video data. We tested on the Sky Time-lapse dataset [1], which has time sequences of sky images. We use the diffusion model pretrained with 256x256 ImageNet data, and then fine tuned on 20 video sequences. Since our method can only handle single channel images like medical images (but should be able to extend to RGB with a little effort), we only use the first channel of the video for fine tuning). The learning rate is set to be the same as CT images, and we observe convergence at 5000 iterations with a batch size of 4. Then we test the inter-slice total-variation (TV) value on the unconditional generation of a 16-frame video (with / without random blending). One key observation is that video data has more coherent frames so it is actually easier for us to train the 3D patch diffusion on it since there is no inter-slice jump which is observed on 3DCT data. We present the z-axis TV values in the table below, which is computed by $||D_z(x)||_1/n$, where $n$ is number of pixels. The ground truth is taken from the test dataset for CT, and an average of 10 videos in sky-timelapse video. The results are based on an average from 10 generated volumes/videos. | Task | Without Blending | With Blending | Ground Truth | |------| ------| ------| ------| |3DCT Generation | 0.0236 | ***0.0035*** | 0.0065 |Skylapse Video Generation | 0.0419 | ***0.0021*** | 0.0036 | We observe with the proposed blending method improves the inter-slice smoothness significantly for both video and 3DCT data. Some video data has smaller inter-slice variation than 3D medical images. We observe that our method can generate 3D images/videos with the inter-slice smoothness comparable to the ground truth. Thus, this prior can also be useful for solving inverse problems in a broad applications. Thanks for the review and reading our rebuttal. Your review is crucial for us to improve our manuscript. Feel free to let us know if there is any remaining question about the manuscript and we will try our best to answer. [1] Zhang, Jiangning, et al. "Dtvnet: Dynamic time-lapse video generation via single still image." Computer Vision–ECCV 2020: 16th European Conference, Glasgow, UK, August 23–28, 2020, Proceedings, Part V 16. Springer International Publishing, 2020.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC