Learning Image Priors through Patch-based Diffusion Models for Solving Inverse Problems

Summary

This work introduces a patch-based diffusion modeling approach to efficiently learn image priors that can be used to solve inverse problems. Particularly the model maintains memory and data efficiency due to the patch-based operating scheme. Experiments are demonstrated in both natural and medical image domains to solve various inverse problems (e.g., CT reconstruction, deblurring, etc.) based on priors learned via patches.

Strengths

- Paper presents a very interesting and novel approach to learn image priors with patch-based diffusion models, and empirically demonstrates rigorous results.

Weaknesses

- Some experimental clarifications are necessary, and writing & quality of figures can be improved.

Questions

- Details regarding how the evaluation metrics PSNR and SSIM are calculated are missing (i.e., in RGB domain, or via the luminance in YCbCr domain)? - The authors' justification on using non-overlapping patches during training is somewhat not clear to me. What exactly is the negative impact of using overlapping patches during training this model? - In Figure 4, results obtained with [16] are surprisingly bad. Is the implementation modified, or did the authors implement themselves? Is it the best case scenario for this method's performance? Also on a different note, did the authors perhaps try using the patch-based diffusion approach from [23] in an unsupervised manner in this simulation? - Considering Algorithm 1's sampling loop, looks like the sampling process goes through all denoising steps without skips (sampling takes T steps)? In that case, perhaps inference time comparisons should also be demonstrated comparatively with other methods. - Visual results (e.g., in Figures 5, 6, 7) should be blended in the compiled PDF not as a PNG/JPEG, but as a high resolution image as a e.g., PDF (more professional looking figure quality needed). Fine differences/details that one should see when zoomed-in can disappear if figures are blended in the paper using a PNG/JPEG type of format, which loses its meaning.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Sufficiently addressed.

Summary

In this work, the authors propose a novel method for learning efficient data priors for entire images by training diffusion models only on image patches. During inference the authors introduce a patch-based position-aware diffusion inverse solver, which obtains the score function for the whole image through scores of individual patches and their positional encoding, using this as the prior for solving inverse problems. Multiple experiments on CT data as well as on the CELEBA dataset are conducted.

Strengths

The idea of including positional information of patches into the diffusive reconstruction process seems novel and promising. Indeed, the authors demonstrate that their proposed approach can compute the score function for entire images without needing to feed the whole image through the network.

Weaknesses

While the patch-based reconstruction for diffusion models appears promising, the evaluation and experiments presented in the paper require further extension to justify publication. Additional evaluations should be conducted against other models, reconstruction methods, inverse problems, and datasets. Furthermore, a sensitivity analysis should be included to examine how the proposed approach performs with different forward operators. Additionally, many of the presented CT reconstructions exhibit significant hallucinations, which is particularly concerning in the context of medical imaging. Overall, my decision is influenced by the weak evaluation of the method.

Questions

I would be interested as to which number of data-samples the patch-based approach outperforms the Vanilla one? How would the proposed approach perform with different forward operators, i.e. different inverse problems on the same datasets? Why do the presented CT reconstructions exhibit such significant hallucinations, and how can this issue be addressed?

Rating

5

Confidence

5

Soundness

3

Presentation

2

Contribution

2

Limitations

Although sparsely addressed in the conclusions, the authors do not give limitations about the proposed approach at all. Especially, the proposed approach seems to have problems concerning hallucinations in the reconstructed CT images.

Summary

This paper proposes a patch-based diffusion model for inverse problems, such as CT reconstruction and natural image deblurring. This method divides images into patches, reducing the size of the input data fed into the network, thereby decreasing memory consumption.

Strengths

This method provides a feasible approach for dividing and merging patches to reduce boundary artifacts for diffusion models.

Weaknesses

1. The innovation is limited. The main contribution of this paper is the application of patch diffusion [1] to inverse problems. However, the experimental results are neither promising nor convincing. [1] Wang, Zhendong, et al. "Patch diffusion: Faster and more data-efficient training of diffusion models." Advances in neural information processing systems 36 (2024). 2. The CT images presented by the author in the paper and appendix are very blurry and not displayed with the correct window level and width, making it impossible to discern imaging details and lacking clinical significance. 3. Even though the CT images are very blurry, it is still evident that all the reconstructed CT images exhibit significant numbers of image artifacts (incorrect organ structures) compared to the ground truth. This is completely unacceptable for medical images. 4. The comparison methods are very limited, with only ADMM-TV, which is a very old method. There are numerous methods for CT reconstruction [2-4], natural image deblurring [5] and super-resolution [6] that the author did not compare with at all. [2] Shen, Liyue, John Pauly, and Lei Xing. "NeRP: implicit neural representation learning with prior embedding for sparsely sampled image reconstruction." IEEE Transactions on Neural Networks and Learning Systems 35.1 (2022): 770-782. [3] Wu, Qing, et al. "Self-supervised coordinate projection network for sparse-view computed tomography." IEEE Transactions on Computational Imaging 9 (2023): 517-529. [4] 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. [5] Tang, Xiaole, et al. "Uncertainty-aware unsupervised image deblurring with deep residual prior." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2023. [6] Wang, Longguang, et al. "Unsupervised degradation representation learning for blind super-resolution." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2021.

Questions

1. I suggest the authors make more improvements for image restoration tasks (inverse problems) to control the reliability of the generated images and reduce artifacts. 2. it is necessary to compare with more SOTA methods to obtain a more comprehensive evaluation.

Rating

5

Confidence

5

Soundness

2

Presentation

2

Contribution

3

Limitations

1. Due to the severe artifacts observed in the visual results (CT reconstruction) and the limited number of comparison methods, the experimental results are hardly convincing. 2. The innovation is very limited. It is based on existing diffusion models and merely applied to the domain of inverse problems (image restoration tasks).

Summary

This manuscript discusses diffusion models for inverse problems. The authors discuss using image patches of the image to improve computational bottlenecks and overcoming the lack of sufficient data in training appropriate surrogate neural network priors for the inversion task. The authors discuss details of their proposed method and illustrate the advantages of their methods on various tasks including CT, deblurring, and superresolution.

Strengths

This manuscript is well-written and structured, making it easy for readers to follow the presented ideas. The authors ground their work in existing literature and reference relevant papers in this field. Data-driven approaches for inverse problems have shown significant advances and this manuscript contributes to this field.

Weaknesses

This work is partly incremental and heavily relies on various previous and cited publication, e.g., [12,18,19]. Furthermore, I assume computational costs for the solution of the inverse problem are extremely large since just a stochastic gradient approach must be utilized in the inversion process (see algorithm 1). These are typically prohibitive for large-scale inverse problems removing the advantages of the learned data-driven prior.

Questions

See Weaknesses.

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes

Summary

The authors present an approach for tile-based training and prediction of diffusion models applied for inverse problem posterior sampling. The core idea is that training is done with random patches and during generation the authors use a differently shifted non-overlapping tiling grid for each iteration of the process. The authors additionally provide the x- and y-pixel-coordinates to the network by encoding them in extra channels and concatenating them. Their approach allows them to generate images without Stichting artefacts. The random tiling during training can be seen as data augmentation and the authors show that this enables their method to be trained on smaller datasets.

Strengths

* Improving the memory requirement for diffusion model is an important problem. Many applications need to process large images in a coherent way, while avoiding stitching artefacts. * I appreciate the fact that the proposed training scheme enables training with less data. And that the authors validate this in an experiment. * The authors show that their method does not depend on the particular smapling scheme or network architecture. I appreciate the generality of the approach.

Weaknesses

* My main criticism is regarding the motivation of the problem. The authors write: "*Directly using overlapping patches would result in sections of the image covered by multiple patches to be updated multiple times, which is inconsistent with the theory of diffusion models.*" The question of how tiling and stitching can be applied for unets, such that the result is equivalent to processing the image as a whole has been explored before. See questions section for details. * The proposed tiling comes at a cost: It limits the range of correlations that can be captured by the diffusion model. This is visible in Figure 4, where the generated images show no stitching artefacts, but also produce nonsensical large scale anatomy. I am missing a discussion of this aspect. In the posterior samples, the effect is not visible, since the input image contains enough information such that long range correlations are not relevant. I believe this would be a problem for inverse problems where the input image contains less information, i.e., when the noise is very severe, or for super resolution with a more extreme resolution factor.

Questions

It is not correct that using overlapping tiles would be "*inconsistent with the theory of diffusion models*". In general, a tiling scheme for unets can be implemented, such that the stitched tiles are identical to the result of processing the image as a whole. This can be achieved by using overlapping tiles (with the correct shift) and disregarding the areas close to the border in the outputs that are influence by padding. I believe a discussion of this can be found in [1]. What would prevent us from applying this approach with a diffusion model in each step? The network could still be trained using patches. While such a tiling approach would cost additional computation time, since it requires overlapping patches, it should theoretically produce guaranteed stitch artefact-free outputs. Do the authors agree, that this well established approach should ideally be a baseline or at least be discussed? The authors could show that their method produces comparable results at reduced computational cost. [1]: Rumberger, Josef Lorenz, et al. "How shift equivariance impacts metric learning for instance segmentation." Proceedings of the IEEE/CVF International Conference on Computer Vision. 2021.

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

I am missing a discussion on the cost of the proposed patch scheme regarding the ability of the network to model long range correlations. See weaknesses.

The rebuttal response

The author's rebuttal addressed some of my concerns, so I have revised my score to borderline accept. However, there are still some unresolved issues, as outlined below. 1. The authors directly addressed the issue of CT artifacts, and I acknowledge their explanation that diffusion-based extremely-limited-angle CT reconstruction inherently results in such artifacts. However, from a medical application perspective, these artifacts are quite concerning. In the rebuttal, the author adds reconstruction results from 60 angles. However, since the selected images primarily show the lungs (which appear as zero in the specified window width), they do not contain enough tissue, unlike abdominal images, to adequately assess the reconstruction quality. 2. In both the paper and the rebuttal, the author points out that performing diffusion on patches can reduce computational costs, enabling to process large-scale images, such as higher resolution and 3D images. In the experiments, both the CT dataset and the CelebA-HQ dataset are 256x256. CelebA-HQ is a standard dataset, but CT data often comes in much higher resolutions, such as 512x512. I suggest the authors conduct experiments on such CT datasets to better support their claims. 3. In summary, while the author's method performs well in terms of quantitative metrics, the inevitable artifacts in the CT data raise some concerns from a medical perspective.

Thanks for the rebuttal - additional questions.

Thank you for the rebuttal. **Regarding the tiling and Stichting method:** I appreciate that you try the suggested stitching mechanism and put results into the pdf. The sampled results look indeed inferior. You write: *"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]. "* I agree that these are the likely reasons for the performance difference. I did not suggest the tiling strategy as an alternative to providing noise level and patch positions as additional inputs (I think these make a lot of sense), but rather as the established method to avoid stitching artefacts, as opposed to the proposed method of using different shifts in each iteration of the generation process. I don't see a reason not to combine the established way of avoiding stitching artefacts in UNets with the additional input information (position and noise level) **Regarding long range correlations:** I believe learning *"different distribution of patches at different locations in the image"* is not the same as learning long range correlations, which means learning which structures on the one side of the image are likely to appear together with structures on the other side of the image. I still don't think it is possible to learn such correlations without looking at the image as a whole.

Response

After considering the author's response, some of my concerns have been alleviated, prompting me to adjust my score to a borderline accept.

response to rebuttal

Thanks to the authors for their rebuttal and detailed comments. Majority of my concerns are answered, and I believe this submission went through a successful rebuttal period. Thus, I also increased my score. Please include the discussions presented here in the revised PDF as well, particularly the inference time comparisons that are provided here. It presents a more fair comparison of the proposed algorithm.