Learning Image Priors Through Patch-Based Diffusion Models for Solving Inverse Problems

We thank the reviewer for their insightful comments.

Comment: Details regarding calculations of PSNR and SSIM are missing

The PSNR and SSIM of RGB images are calculated in the RGB domain. Data preprocessing consisting of dividing all the RGB values by 255 was done first, so all the reconstructed images have values between 0 and 1. The PSNR and SSIM values were then computed from these images.

Comment: Justification for using non-overlapping patches is unclear

Using overlapping patches during training would significantly increase the computational cost. Section A.6 details the theory behind which the original score matching process (that must be done on the entire image), used to train the network, can be reduced to score matching on individual patches. In particular, in going from equation (A.6) to equations (A.7) and (A.8), the assumption that the patches do not overlap is made to bring the sum out of the norm. If overlapping patches were to be used, it would be necessary to backpropagate through all the terms of the sum during the training. On the other hand, when nonoverlapping patches are used, we can perform score matching on individual patches and only the loss of these individual patches needs to be backpropagated through the network.

Comment: Results obtained by [16] are surprisingly bad

In Figure 4, for [16], we used the code shared by the paper’s authors. However, there is a key difference between the figures we generated and the figures generated in [16]. In [16], for the best results, some portion of the training time must be spent on learning the distribution of the entire image (without patches). Then during generation, the entire image is used as an input to the network. However, the goal for our paper was to avoid needing to input the entire image into the network both at training time and generation. Therefore, when running [16] in Figure 4, we trained the network only on patches of images and we generated full size images by first generating patches of the images (with positional encoding information) and then simply stitching them together.

Comment: Did the authors try using [23] in an unsupervised manner

When using [23] in an unsupervised manner, it is necessary to first train an unsupervised diffusion model (the patch-based networks from our work suffices for this) and then apply the unsupervised network to solve the inverse problem. The original paper is able to use the conditional network to enforce data consistency with the measurement, but with an unsupervised network, it is necessary to add in an additional step in the inference loop that enforces data consistency. We ran experiments using DPS as the data consistency strategy that is consistent with PaDIS. Further note that [23] has an additional tunable parameter which is the amount of overlap between patches.

Table 6 shows the results of this method when this parameter has been tuned under the name Patch Averaging. The table shows the method can obtain reasonable results but is outperformed by our proposed method. Additionally, the optimal overlap parameter value requires a significant amount of overlap between patches (approximately 1/4 of the patch dimension in both x and y) which increases the number of patches per iteration whose score function must be evaluated using the network. Finally, while the empirical results are reasonable, there is a lack of mathematical justification in [23] for the procedure of averaging the predicted noises of overlapping patches and future work is required to theoretically justify (from a probability distribution perspective) this method. In the revision, we will include more visual examples of this method compared with the others.

Comment: Inference time comparisons should be demonstrated with other methods

Algorithm 1 indeed indicates that none of the sampling steps are skipped. For a fair comparison with other diffusion model based methods in Table 5, we used the same number of steps (1000) for all the methods; an increase in image quality was present with an increased number of steps for all the methods. DPS requires more time per image due to the computation of the gradient of the norm term which involves backpropagating through the network, and predictor-corrector sampling involves two network evaluations per iterations. The average reconstruction time per image in seconds for the methods in Tables 1 and 5 are shown below for 20 view CT.

Baseline: 0.1

ADMM-TV: 0.7

Whole image diffusion: 172

PaDIS (VE-DPS): 195

Langevin dynamics: 98

Predictor-corrector: 189

VE-DDNM: 105

Comment: Visual results should be presented as a high resolution image

The global rebuttal PDF page has some examples of higher resolution images for different CT reconstruction experiments including 60 view reconstruction and fan beam CT. These images are displayed with higher contrast and should be more helpful in performing clinical diagnosis.

We thank the reviewer for their insightful comments.

Comment: Additional evaluations should be conducted against other inverse problems, sensitivity analysis should be included with different forward operators

We conducted more experiments with different forward models: namely 60 view parallel beam CT, 180 view fan beam CT, and deblurring with a larger kernel of size 19x19. The results are shown in Table 7 and further demonstrate that our proposed method outperforms various SOTA methods for a large variety of forward models. Additionally, Table 6 contains a comparison with a wider variety of inverse problem solving methods. These comparisons with several other SOTA methods strengthen the evaluation of our method.

Comment: CT reconstructions exhibit significant hallucinations

The authors acknowledge that the images obtained by the generative models investigated including the proposed method for 20 view CT reconstruction show some hallucinations and artifacts. This is a natural consequence of using extreme compressed sensing with ultra-sparse views: normally, to reconstruct a 256x256 image requires (pi/2*256)=402 views, so for the 20 view experiments, the measurements have been compressed by a factor of 20. Due to this lack of information, it is very hard for any model to perform a diagnostic-quality reconstruction, though our proposed method (and the other diffusion model methods, to a lesser extent) are able to partially fill in this information through learning a strong image prior. The alternative methods that do not learn a prior perform significantly worse in terms of the shown metrics and exhibit severe blurring and artifacts. In clinical settings, it is much more common to perform patient diagnosis with CT scans consisting of hundreds of views. To illustrate this point, we perform experiments with 60 view CT, where our proposed method is able to obtain excellent quality images as shown in Figure B.1: essentially no artifacts are visible. (We show the potential of our proposed method to reconstruct images with ultra-sparse views with a decent image quality, which can be potentially used for other clinical applications such as patient positioning.)

Comment: What is the number of samples for which the proposed method is better

When looking at PSNR, the number of samples out of the 25 test samples in which the patch-based approach outperformed the vanilla one is as follows: 23 for 20 view CT, 25 for 8 view CT, 20 for deblurring, 16 for superresolution. We will add this information to the supplement.

Comment: Authors do not give limitations of the proposed method

One limitation of the proposed approach (and all diffusion approaches) is that they tend to be slower than optimization based approaches, plug and play methods, and model-based learning methods. Another limitation of generative modeling approaches is the potential to hallucinate, particularly when the measurements are very compressed and contain little information. This is a limitation of most generative models, as illustrated by the visual examples from the whole image diffusion model for sparse view CT. This problem can be resolved by obtaining more projection views in a CT scan, depending on the application needs.

We thank the reviewer for their insightful comments.

Comment: Innovation is limited

The main innovation of our method is a method to formulate a diffusion based image prior from solely the patches of the image. Diffusion models are known for requiring a large amount of memory for training and inference and extending them for large scale images is a challenging problem. This work illustrates how image priors for very large images can be learned by learning priors of patches, a much less memory intensive task. This is unlike the work of [1], which ultimately still requires training on whole images as well as inputting the whole image into the network to generate images. We also demonstrate in Table 3 that our proposed method can learn a reasonable prior for dataset sizes much smaller than is normally used for training diffusion models and the advantage over whole image diffusion models becomes more pronounced for small datasets. Finally, unlike previous patch-based diffusion model papers such as [16] and [17] that can only be used for generating images, we show how our method learns a full image prior from only image patch training that can be coupled with most diffusion inverse problem solving algorithms, which is not a trivial extension based on [16] and [17].

Comment: CT images are blurry and displayed with the wrong window level

We displayed CT images corresponding to our new CT experiments in Figure B.1 using a narrower window size of 800 to 1200 HU and with higher resolution. These images show better contrast between different organs and are more useful for obtaining a clinical diagnosis. In the revision, we will redisplay all the CT reconstruction images with this higher level of contrast.

Comment: CT images contain significant number of image artifacts

The authors acknowledge that the images obtained by the generative models investigated including the proposed method for 20 view CT reconstruction show some hallucinations and artifacts. This is a natural consequence of using extreme compressed sensing with ultra-sparse views: normally, to reconstruct a 256x256 image requires (pi/2*256)=402 views, so for the 20 view experiments, the measurements have been compressed by a factor of 20. Due to this lack of information, it is very hard for any model to perform a diagnostic-quality reconstruction, though our proposed method (and the other diffusion model methods, to a lesser extent) are able to partially fill in this information through learning a strong image prior. The alternative methods that do not learn a prior perform significantly worse in terms of the shown metrics and exhibit severe blurring and artifacts. In clinical settings, it is much more common to perform patient diagnosis with CT scans consisting of hundreds of views. To illustrate this point, we perform experiments with 60 view CT, where our proposed method is able to obtain excellent quality images as shown in Figure B.1: essentially no artifacts are visible. (We show the potential of our proposed method to reconstruct images with ultra-sparse views with a decent image quality, which can be potentially used for other clinical applications such as patient positioning.)

Comment: Comparison methods are very limited

[4] is a method that applies 2D diffusion models to solve 3D inverse problems including CT reconstruction, whereas our proposed method applies for 2D inverse problems, so it cannot directly be applied. However, the sampling algorithm that is used is the predictor-corrector sampling algorithm, which we initially compared to in Table 5 and now added a more comprehensive comparison in Table 6. This sampler did not perform as well as the one chosen for PaDIS (DPS) by the quantitative metrics.

[2] and [3] are self-supervised methods for CT reconstruction, which means that network is trained during reconstruction time, substantially slowing down the algorithm. [5] is a deep image prior method that also requires network training at inference time. Furthermore, [2], [3], [5], and [6] are all problem specific methods for which a generalization to the other types of inverse problems would be nontrivial. Our proposed method, along with most of the methods we compared to in Table 6, are easily generalizable methods that can solve a wide variety of inverse problems. Furthermore, the training process for our algorithm need only happen once per dataset, and no network training is required during reconstruction time. Due to these fundamental differences between our method and [2], [3], [5], [6], we believe that comparisons with those methods would not be fair.

Nevertheless, to provide a more complete evaluation of our method, we included additional comparisons with plug and play (PnP) methods and other diffusion inverse solvers in Table 6. These methods share the similarity with our method in that network training only needs to be done once per dataset, and the same trained network can be used for different types of inverse problems, allowing for greater flexibility. Table 6 consists of an expanded comparison between various methods. We implemented various diffusion inverse solving methods [1], [7], [19] in conjunction with our patch-based prior. We included two additional patch based methods from [23] and [69] where we applied [23] in an unsupervised way by using the same unsupervised network trained in our proposed method and adding a DPS step during reconstruction. We also implemented two plug and play (PnP) methods by first training denoisers on CT images and the CelebA dataset and then applying these denoisers in an unsupervised way to solve the inverse problems. Optimal hyperparameters for all these methods were found through searching.

In all cases, our proposed method outperformed the comparison methods. In the revision, we will add visual examples of these methods. These comparisons with several other SOTA methods strengthen the evaluation of our method.

We thank the reviewer for their insightful comments.

Comment: Work is partly incremental, heavily relies on cited publications

The papers [12] and [18] apply diffusion models to solve 3D reconstruction problems, whereas the proposed method performs experiments on 2D reconstruction problems using a patch-based prior. [19] uses the predictor-corrector method for solving medical imaging problems, which we compared to in Table 5. We revised the labels in Table 5 to make this more clear. The method of diffusion inverse problem solving most closely related to our method is DPS [5]. However, the most significant contribution of the proposed method is a patch-based image prior that requires only patch inputs to a neural network and can be paired with any diffusion inverse solving algorithm, as illustrated in Table 5.

Comment: Computational costs are large and requires a stochastic gradient approach

We acknowledge that the computational costs of the proposed method will exceed that of more traditional reconstruction methods, as is true for almost all diffusion based methods, as a tradeoff for achieving better image quality. However, the computational cost of the proposed method is similar to other diffusion based methods. The average reconstruction time per image in seconds for the methods in Tables 1 and 5 are shown below for 20 view CT. Notably, our proposed method takes only slightly more time than the approach using diffusion models trained on entire images while greatly reducing the memory needed and improving the result. We will add this table to the supplement.

Baseline: 0.1

ADMM-TV: 0.7

Whole image diffusion: 172

PaDIS (VE-DPS): 195

Langevin dynamics: 98

Predictor-corrector: 189

VE-DDNM: 105

The stochastic approach taken in Algorithm 1 allows the runtime of the algorithm to be similar to other diffusion methods while eliminating artifacts that would otherwise persist between boundaries of patches. The approach is similar to the one taken in the paper below: instead of computing the score function multiple times each iteration, stochastically choose one of them to compute each iteration. Over the course of hundreds of iterations throughout the reconstruction process, the stochastic approximation of the score function becomes more accurate. Furthermore, this approach does not sacrifice the advantages of using this data-driven prior, as Figure 4 shows that our proposed method can still be used to unconditionally generate fairly realistic images.

S. Lee, H. Chung, M. Park, J. Park, W.-S. Ryu, and J. C. Ye. “Improving 3D imaging with pre-trained perpendicular 2D diffusion models”. In: Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023, pp. 10710–10720.

We thank the reviewer for their insightful comments.

Comment: Tiling scheme for unets can be implemented with overlapping patches

We implemented the method provided in the above reference [1] while using the same trained network, hyperparameters, and DPS for inverse problem solving. The network consisted of 3 layers of upsampling/downsampling and each layer involved pooling with a factor of 2. The largest patches that the network was trained on had size 56. To satisfy the hypotheses of the paper, the overlap between patches was to be a multiple of 8; we set it equal to 8 to minimize the number of patches for image partitioning. Hence each 256x256 image was divided into 36 overlapping patches, compared to 25 for PaDIS.

Table 6 shows the results of using this method under the name Patching Stitching. For the shown inverse problems, the method obtains reasonable results but performs worse than our proposed method. Despite this, the reconstructed images did not appear to exhibit any boundary artifacts along patch boundaries, showing that the method of [1] worked in that aspect. The increased number of patches necessary for that approach increased the runtime of the reconstruction algorithm by approximately 30%.

We also examined the unconditionally generated images using the method of [1] which are shown in Figure B.2. Although these images exhibited relatively smooth features without any clear boundary artifacts between patches (unlike the clear artifacts visible in the middle row of Figure 4 generated by naive patch stitching), the overall structure was highly inconsistent with the CT images in the dataset. This is in contrast with the images generated by PaDIS as shown in the bottom row of Figure 4. Therefore, although [1] can result in smooth images, the lack of patch shifting means that the learned image prior differs from the one in Eq. (3), resulting in unrealistic looking generated images, which indicated the underlying data distribution cannot be well captured. The significantly worse prior learned by this method is reflected in the worse results for inverse problem solving and especially when the measurements are highly compressed.

There are two aspects of the UNet used in our application that likely caused that patch stitching method to fail to learn the prior well. Firstly, our diffusion UNet takes in an additional scalar input indicating the noise level of the noisy image being input into the network. This scalar input is processed through a sinusoidal positional encoding before being embedded into the layers of the UNet through an attention mechanism. Secondly, our network takes in the positional embedding of the location of the patch via concatenation along the channel dimension. Thus, our network learns the score function of patches while also incorporating the location of the patch, making it different from traditional UNets used for segmentation [1].

Comment: Patches limit the ability of the network to learn long range correlations

Learning longer range correlations within an image is assisted by the method of using positional encoding of patches as an input to the network: it allows our network to learn a different distribution of patches at different locations in the image. Thus, provided that the location of the central object to be imaged is in a relatively consistent location (as is the case for CT scans or human faces), the network can learn that, for instance, the spine of the CT scan is typically around the middle bottom section of the image. Such learning is consistent with the generated results in Figure 4. Nevertheless, generating whole images with realistic large scale anatomy is challenging, and we show the generation results to demonstrate that they appear somewhat reasonable while emphasizing the focus on solving inverse problems.

Figure 5 contains some examples of reconstructed images from 8 view CT, a very compressed sensing problem. Normally, to reconstruct a 256x256 image would require (pi/2*256)=402 views, so the compression is a factor of 50. In this case, the images show reasonable large scale anatomy.