Patch-Based Diffusion Models Beat Whole-Image Models for Mismatched Distribution Inverse Problems

## Abstract

Abstract Diffusion models have achieved excellent success in solving inverse problems
due to their ability to learn strong image priors,
but existing approaches require a large training dataset of images
that should come from the same distribution as the test dataset.
When the training and test distributions are mismatched,
artifacts and hallucinations can occur in reconstructed images due to the incorrect priors.
In this work, we systematically study out of distribution (OOD) problems where a known training distribution is first provided.
We first study the setting where only a single measurement obtained from the unknown test distribution is available.
Next we study the setting where a very small sample of data belonging to the test distribution
is available, and our goal is still to reconstruct an image from a measurement that came from the test distribution.
In both settings, we use a patch-based diffusion prior
that learns the image distribution solely from patches.
Furthermore, in the first setting, we include a self-supervised loss
that helps the network output maintain consistency with the measurement.
Extensive experiments show that in both settings,
the patch-based method can obtain high quality image reconstructions
that can outperform whole-image models
and can compete with methods that have access to large in-distribution training datasets.
Furthermore, we show how whole-image models are prone to memorization and overfitting,
leading to artifacts in the reconstructions, while a patch-based model can resolve these issues.

## 1 Introduction

In image processing, inverse problems are of paramount importance
and consist of reconstructing a latent image ${\bm{x}}$ from a measurement
${\bm{y}}=\mathcal{A}({\bm{x}})+\bm{\varepsilon}$.
Here, $\mathcal{A}$ represents a forward operator
and $\bm{\varepsilon}$ represents random unknown noise.
By Bayes’ rule, $\log p({\bm{x}}|{\bm{y}})$ is proportional to $\log p({\bm{x}})+\log p({\bm{y}}|{\bm{x}})$,
so obtaining a good prior $p({\bm{x}})$ is crucial for recovering ${\bm{x}}$
when ${\bm{y}}$ contains far less information than ${\bm{x}}$.
Diffusion models obtain state-of-the-art results for
learning a strong prior and sampling from it,
so similarly competitive results can be obtained when using them to solve inverse problems (Chung et al., ; ; Song et al., ; Wang et al., ; Kawar et al., ; Li et al., ).

However, training diffusion models well requires vast amounts of clean training data
(Song et al., ; Ho et al., ),
which is infeasible to collect in many applications such as medical imaging
(Chung et al., ; Song et al., ; Jalal et al., ),
black hole imaging (Feng et al., ; ),
and phase retrieval (Li et al., ; Wu et al., ).
In particular, for very challenging inverse problems
such as black hole imaging (Feng et al., )
and Fresnel phase retrieval (Gureyev et al., ),
no ground truth images are known and one only has a single measurement $\bm{y}$ available.
In other applications such as dynamic CT reconstruction (Reed et al., )
and single photo emission CT (Li et al., ),
obtaining a high quality measurement that can lead
to a reconstruction that closely approximates the ground truth
can be slow or potentially harmful to the patient,
so only a very small dataset of clean images are available.
Thus, in this paper we consider two settings:
the single measurement setting in which we are given one measurement $\bm{y}$ whose corresponding $\bm{x}$ belongs to a different distribution from the training dataset,
and the small dataset setting in which we are only given a small number of samples $\bm{x}$ that belong to the same distribution as the test dataset.

In recent years, some works have aimed to address these problems
by demonstrating that diffusion models have a stronger generalization ability
than other deep learning methods (Jalal et al., ),
so slight distribution mismatches between the training data and test data
may not significantly degrade the reconstructed image quality.
However, in cases of particularly compressed or noisy measurements,
as well as when the test data is severely out of distribution (OOD) with a significant domain shift,
an improper choice of training data leads to an incorrect prior
that causes substantial image degradation and hallucinations (Feng et al., ; Barbano et al., ).
To address these challenges in the single measurement case,
recent works use each measurement $\bm{y}$ to adjust the weights of a diffusion network at reconstruction time
Barbano et al. (); Chung & Ye (),
aiming to shift the underlying prior learned by the network
toward the appropriate prior for the latent image in the test case.
However, as the networks have huge numbers of weights,
an intricate and parameter-sensitive refining process of the network is required
during reconstruction to avoid overfitting to the measurement.
Furthermore, there is still a substantial gap in performance between methods
using an OOD prior and methods using an in-distribution prior.
Finally, these methods have only been tested in medical imaging applications (Barbano et al., ; Chung & Ye, ).
On the other hand, in the small dataset case,
various methods (Moon et al., ; Zhang et al., ) have been devised
to fine-tune a diffusion model on an OOD dataset,
but these methods still require several hundred images
and have not used the fine-tuned network to solve inverse problems.

Patch-based diffusion models have shown success both for image generation
(Wang et al., ; Ding et al., )
and for inverse problem solving (Hu et al., ).
In particular, the method of Hu et al. () involves training networks
that take in only patches of images at training and reconstruction time,
learning priors of the entire images from only priors of patches.
In cases of limited data,
Hu et al. () shows that patch-based diffusion models
outperform whole image models for solving certain inverse problems.
These works motivate our key insight that patch-based diffusion priors
potentially obtain stronger generalizability than whole-image diffusion priors
for both the single measurement setting and the small dataset setting due to a severe lack of data.
Inspired by this, we propose to utilize patch-based diffusion models
to tackle the challenges arising from mismatched distributions and lack of data in a unified way.
We first develop a method to take a network trained on patches of a mismatched distribution
and adjust it on the fly in the single measurement setting.
We also show how in the small dataset setting, fine-tuning a patch-based network
results in a much better prior than fine-tuning a whole-image network,
leading to higher quality reconstructed images.

In summary, our contributions are as follows:

- •
We integrate the patch-based diffusion model framework
with the deep image prior (DIP) framework
to correct for mismatched distributions in the single measurement setting.
Experimentally, we find this approach beats using whole-image models
in terms of quantitative metrics and visual image quality in image reconstruction tasks,
as well as achieving competitive results with methods using in-distribution diffusion models.
- •
We show that in the small dataset setting,
fine-tuning patch-based diffusion models is much more robust than whole-image models
and very little data is required to obtain a reasonable prior for solving inverse problems.
- •
We demonstrate experimentally that when fine-tuning on very small datasets,
whole image diffusion models are prone to overfitting and memorization,
which severely degrades reconstructed images,
while patch-based models are much less sensitive to this problem.

## 2 Background and Related Work

Diffusion models and inverse problems.
In a general framework,
diffusion models involve the forward stochastic differential equation (SDE)

$$ $\mathrm{d}{\bm{x}}=-\frac{\beta(t)}{2}\,{\bm{x}}\,\mathrm{d}t+\sqrt{\beta(t)} \,\mathrm{d}{\bm{w}},$ (1) $$

where $t\in[0,T]$, ${\bm{x}}(t)\in\mathbb{R}^{d}$,
and $\beta(t)$ is the noise variance schedule of the process.
This process adds noise to a clean image and ends with an image indistinguishable from Gaussian noise  (Song et al., ).
Thus, the distribution of ${\bm{x}}(0)$ is the data distribution
and the distribution of ${\bm{x}}(T)$ is (approximately) a standard Gaussian.
Then the reverse SDE has the form (Anderson, ):

$$ $\mathrm{d}{\bm{x}}=\left(-\frac{\beta(t)}{2}-\beta(t)\nabla_{{\bm{x}}_{t}}\log p _{t}({\bm{x}}_{t})\right)\,\mathrm{d}t+\sqrt{\beta(t)}\,\mathrm{d}\overline{{ \bm{w}}}.$ (2) $$

Score-based diffusion models involve training a neural network to learn the score function
$\nabla_{{\bm{x}}_{t}}\log p_{t}({\bm{x}}_{t})$,
from which one can start with noise and run the reverse SDE
to obtain samples from the learned data distribution.

When solving inverse problems, it is necessary to instead sample from $p({\bm{x}}_{T}|{\bm{y}})$,
so the reverse SDE becomes

$$ $\mathrm{d}{\bm{x}}=\left(-\frac{\beta(t)}{2}-\beta(t)\nabla_{{\bm{x}}_{t}}\log p _{t}({\bm{x}}_{t}|{\bm{y}})\right)\,\mathrm{d}t+\sqrt{\beta(t)}\,\mathrm{d} \overline{{\bm{w}}}.$ (3) $$

Unfortunately, the term $\log p_{t}({\bm{x}}_{t}|{\bm{y}})$ is difficult to compute
from the unconditional score $\nabla_{{\bm{x}}_{t}}\log p_{t}({\bm{x}}_{t})$ alone.
Liu et al. (), Chung et al. (), and Ozdenizci & Legenstein () among others
proposed directly learning this conditional score $\nabla_{{\bm{x}}_{t}}\log p_{t}({\bm{x}}_{t}|{\bm{y}})$ instead.
However, this process requires paired data $({\bm{x}},{\bm{y}})$
between the image domain and measurement domain for training,
instead of just clean image data.
Furthermore, the learned conditional score function
is suitable only for the particular inverse problem for which it was trained,
limiting its flexibility.

For greater generalizability,
it is desirable to apply the unconditional score $\nabla_{{\bm{x}}_{t}}\log p_{t}({\bm{x}}_{t})$
to be able to solve a wide variety of inverse problems.
Thus, many works
have been proposed to approximate the conditional score in terms of the unconditional one
(Wang et al., ; Chung et al., ; ; Kawar et al., ).
Notably, Peng et al. ()
unified various diffusion inverse solvers (DIS) into two categories:
the first consists of direct approximations to $p_{t}({\bm{y}}|{\bm{x}}_{t})$,
and the second consists of first approximating $\mathbb{E}[{\bm{x}}_{0}|{\bm{x}}_{t},{\bm{y}}]$
(typically through an optimization problem balancing the prior and measurement)
and then applying Tweedie’s formula (Efron, ) to obtain

$$ $\nabla\log p_{t}({\bm{x}}_{t}|{\bm{y}})=\frac{\mathbb{E}[{\bm{x}}_{0}|{\bm{x}} _{t},{\bm{y}}]-{\bm{x}}_{t}}{\sigma_{t}^{2}},$ (4) $$

where $\sigma_{t}$ is the noise level of ${\bm{x}}_{t}$.
All of these methods require a large quantity of clean training data
that should come from the distribution $p({\bm{x}})$
whose score is to be learned, which may not be available in practice.

Methods without clean training data.
When no in-distribution data is available, one approach is to use traditional methods
that do not require any training data,
such as total variation (TV) (Li et al., )
or wavelet transform (Daubechies, ) regularizers
that encourage image sparsity.
More recently, plug and play (PnP) methods have risen in popularity
(Sun et al., ; Sreehari et al., ; Hong et al., ; );
these methods use a denoiser to solve general inverse problems.
Although these methods often use a trained denoiser,
Ryu et al. ()
found that using an off-the-shelf denoiser such as block matching 3D (Dabov et al., )
can yield competitive results.
Nevertheless, with the rise of deep learning in image processing applications,
methods that harness the power of these tools may be desirable.

The deep image prior (DIP) is an extensively studied self-supervised method
that is popular when no training data is available
and reconstruction from a single meausurement $\bm{y}$ is desired.
The method consists of training a network $f_{\bm{\theta}}$ using the loss function

$$ $L(\bm{\theta})=\|{\bm{y}}-\mathcal{A}(f_{\bm{\theta}}({\bm{z}}))\|_{2}^{2}, \quad{\bm{z}}\sim\mathcal{N}(0,{\bm{I}}),$ (5) $$

so that $f_{\bm{\theta}}({\bm{z}})$ produces the reconstruction.
Although the neural network acts as an implicit regularizer
whose output tends to lie in the manifold of clean images,
DIP is prone to overfitting (Ulyanov et al., ).
Various methods have been proposed
involving early stopping, regularization, and network initialization
(Liu et al., ; Jo et al., ; Barbano et al., ).
Nevertheless, the method is very sensitive to parameter selection and implementation
and can take a long time to train (Jo et al., ).

Most DIS methods learn a prior from a large collection of clean in-distribution training images,
but recently Barbano et al. () and Chung & Ye ()
proposed self-supervised diffusion model methods that are based off the DIP framework.
These methods involve alternating between the usual reverse diffusion update step
to gradually denoise the image
and a network refining step in which the score network parameters are updated via the loss function

$$ $L(\bm{\theta})=\|{\bm{y}}-\mathcal{A}(\text{CG}(\hat{{\bm{x}}}_{0|t}({\bm{x}}_ {t};\bm{\theta})))\|_{2}^{2}$ (6) $$

where conjugate gradient (CG) descent is used to enforce data fidelity.
This CG step consists of solving an optimization of the form

$$ $\operatorname*{arg\,min}_{{\bm{x}}}\frac{\gamma}{2}\|{\bm{y}}-\mathcal{A}({\bm {x}})\|_{2}^{2}+\frac{1}{2}\|{\bm{x}}-\hat{{\bm{x}}}_{0|t}\|_{2}^{2},$ (7) $$

where $\gamma$ is a tradeoff parameter controlling the strength of the prior versus the measurement.
Crucially, these methods introduce an additional LoRA module (Hu et al., ) to the network
and the original network parameters are frozen when backpropagating the loss,
which helps to avoid overfitting
the whole-image model.
Nevertheless, many technical tricks are required (Chung & Ye, )
involving noisy initializations and early stopping to obtain good results and avoid artifacts.
Our patch-based model avoids this overfitting issue.

Diffusion model fine-tuning.
In the small dataset setting,
various fine-tuning methods exist
to shift the underlying prior learned by a score network
away from a mismatched distribution and toward a target distribution.
Given a pretrained diffusion network on a mismatched distribution,
Moon et al. (), Zhang et al. (), and Zhu et al. () among others
have studied ways to fine-tune the network to the desired dataset.
These methods generally involve freezing certain layers of the original network,
appending extra modules that contain relatively few weights,
or modifying the loss function to capture details that differ greatly between distributions.
However, these methods usually still require thousands of images from the desired distribution
and focus on image generation.
When solving inverse problems,
the reconstructed image should be consistent with the measurement $\bm{y}$, reducing the number of degrees of the freedom for the image compared to generation,
so with proper fine-tuning the data requirement should be lower.

## 3 Methods

### 3.1 Patch-based prior

Figure: Figure 1: Schematic for zero padding and partitioning image into patches. Each index $i$ represents one of the $M^{2}$ possible ways to choose a patch location.
Refer to caption: x1.png

We adapt the patch-based diffusion model framework of Hu et al. ();
we zero pad the image by an amount $P$ on each side
and analyze the distribution of the resulting image ${\bm{x}}$.
Then we assume the true underlying data distribution takes the form

$$ $p({\bm{x}})=\prod_{i=1}^{M^{2}}p_{i,B}({\bm{x}}_{i,B})\prod_{r=1}^{(k+1)^{2}}p _{i,r}({\bm{x}}_{i,r})/Z,$ (8) $$

where ${\bm{x}}_{i,B}$ represents the aforementioned bordering region of ${\bm{x}}$
that depends on the specific value of $i$,
$p_{i,B}$ is the probability distribution of that region,
${\bm{x}}_{i,r}$ is the $r$th $P\times P$ patch
when using the partitioning scheme corresponding to the value of $i$,
$p_{i,r}$ is the probability distribution of that region,
and $Z$ is an appropriate scaling factor.

For training, we use a neural network $D_{\bm{\theta}}({\bm{x}},\sigma_{t})$
that accepts a noisy image ${\bm{x}}$ and the noise level $\sigma_{t}$.
For each patch, we define the $x$ positional array
as the 2D array consisting of the $x$ positions of each pixel of the image scaled between -1 and 1.
To allow the network to learn different patch distributions at different locations in the image,
we extract the corresponding patches of these positional arrays
and concatenate them along the channel dimension of the noisy image patch
and treat the entire array as the network input.
Since we are using a patch-based prior, we perform denoising score matching on patches of an image instead of the whole image. Hence, the training loss is given by

$$ $\operatorname*{arg\,min}_{\bm{\theta}}\mathbb{E}_{t\sim\mathcal{U}(0,T)} \mathbb{E}_{{\bm{x}}\sim p({\bm{x}})}\mathbb{E}_{\bm{\varepsilon}\sim\mathcal{ N}(0,\sigma_{t}^{2}I)}\|D_{\bm{\theta}}({\bm{x}}+\bm{\varepsilon},\sigma_{t})- {\bm{x}}\|_{2}^{2},$ (9) $$

where ${\bm{x}}\sim p({\bm{x}})$ represents a patch drawn from a sample of the training dataset, $\sigma_{t}$ is a predetermined noise schedule, and $\mathcal{U}$ represents the uniform distribution.

### 3.2 Single measurement setting

Consider the first case where only the measurement $\bm{y}$ is given,
and no in-sample training data is available.
For each specific measurement $\bm{y}$,
the DIP framework optimizes
the network parameters $\bm{\theta}$ via the self-supervised loss
(5)
from the predicted reconstructed image.
Diffusion models provide a prediction of the reconstructed image at each timestep:
namely, the expectation of the clean image $\mathbb{E}[{\bm{x}}_{0}|{\bm{x}}_{t}]$
is approximated by the denoiser $D_{\bm{\theta}}({\bm{x}}_{t})$ via Tweedie’s formula.
Then the expectation conditioned on the measurement
$\mathbb{E}[{\bm{x}}_{0}|{\bm{x}}_{t},{\bm{y}}]$ can be obtained
through one of many methods of enforcing the data fidelity constraint.

We begin with the unconditional expectation by leveraging the patch-based prior.
Following (8), we apply Tweedie’s formula
to express the denoiser of ${\bm{x}}$ in terms of solely the denoisers of the patches of ${\bm{x}}$.
Because the outermost product is computationally very expensive,
in practice we approximate $D_{\bm{\theta}}({\bm{x}})$
using only a single randomly selected value of $i$
for each denoiser evaluation:

$$ $D_{\bm{\theta}}({\bm{x}})\approx D_{i,B}({\bm{x}}_{i,B})+\sum_{r=1}^{(k+1)^{2} }D_{i,r}({\bm{x}}_{i,r}).$ (10) $$

By definition, $D_{i,B}({\bm{x}}_{i,B})=0$
and we compute each $D_{i,r}({\bm{x}}_{i,r})$ with the network.
Note that (10)
provides an unconditional estimate of the clean image;
to obtain a conditional estimate $D_{\bm{\theta}}({\bm{x}}_{t}|{\bm{y}})$ of the clean image,
we run $M$ iterations of the conjugate gradient descent algorithm
for minimizing
$\|{\bm{A}}\bm{x}-\bm{y}\|_{2}$,
initialized with the unconditional estimate (Chung et al., ).

The image that is being reconstructed
might not come from the distribution of the training images.
Hence, the estimate $D_{\bm{\theta}}({\bm{x}}_{t}|{\bm{y}})$ may be far from the true denoised image.
Thus, we use $\bm{y}$ to update the parameters of the network
in a way such that $D_{\bm{\theta}}({\bm{x}}_{t}|{\bm{y}})$ becomes more consistent with the measurement:

$$ $\bm{\theta}\leftarrow\operatorname*{arg\,min}_{\bm{\theta}}\|\bm{y}-{\bm{A}}\, D_{\bm{\theta}}(\bm{x}_{t}|\bm{y})\|_{2}^{2}.$ (11) $$

Figure: Algorithm 1 Single Measurement Inverse Solver

Previously, additional LoRA parameters (Hu et al., ) have been used
as an injection to the network
to leave the original parameters unchanged during this process (Barbano et al., ; Chung & Ye, ).
However, the effect of using different ranks for LoRA
versus other methods of network fine-tuning on DIS has not been studied extensively,
so we opt to update all the weights of the network in this step.
Appendix A.3 shows results from using the LoRA module.

Crucially, iterative usage of CG for computing the conditional denoiser
allows for simple and efficient backpropagation through this loss function,
a task that would be much more computationally challenging
if another DIS such as Chung et al. () or Wang et al. () were used.
Furthermore, because the number of diffusion steps is large
and the change in ${\bm{x}}_{t}$ is small between consecutive timesteps,
we apply this network refining step
only for certain iterations of the diffusion process,
reducing the computational burden.

After this step,
we apply the refined network to compute a new estimate of the score of ${\bm{x}}_{t}$
and then use it to update ${\bm{x}}_{t}$.
Similar to the network refining step,
we use the stochastic version of the denoiser given by (10)
rather than the full version.
Hu et al. () showed that for patch-based priors,
Langevin dynamics Song & Ermon () works particularly well as a sampling algorithm,
so we use it here in conjunction with CG steps to enforce data fidelity.
Algorithm 1
summarizes the entire method
for cases where only a single measurement $\bm{y}$ is available.

### 3.3 Small dataset setting

Now turn to the case
where we have trained a diffusion model on OOD data,
but we also have a very small dataset of in-distribution test data
that we can use to fine-tune the model.
When fine-tuning,
we initialize the network with the checkpoint trained on OOD data
and then use the denoising score matching loss function
to fine-tune the network on in-distribution data.
Wang et al. () found that improved image generation performance can be obtained
by training with varying patch sizes,
as opposed to fixing the patch size to the one used during inference.
Here, we apply a varying patch size scheme during fine-tuning also as a method of data augmentation.
We use the UNet architecture in Ho et al. ()
that can accept images of different sizes.
Hence, the loss becomes

$$ $\operatorname*{arg\,min}_{\bm{\theta}}\mathbb{E}_{t\sim\mathcal{U}(0,T)} \mathbb{E}_{{\bm{x}}\sim p_{d}({\bm{x}})}\mathbb{E}_{\bm{\varepsilon}\sim \mathcal{N}(0,\sigma_{t}^{2}I)}\|D_{\bm{\theta}}({\bm{x}}+\bm{\varepsilon}, \sigma_{t})-{\bm{x}}\|_{2}^{2},$ (12) $$

where ${\bm{x}}\sim p_{d}({\bm{x}})$ represents the drawing a randomly sized patch from an image belonging to the fine-tuning dataset.
Appendix A.5 provides full details of the training process.

At reconstruction time, we assume that our network has been fine-tuned reasonably to our dataset.
Thus, we remove the network refining step in Algorithm 1
and keep the weights fixed throughout the entire process.
We still use the same CG descent algorithm to enforce data fidelity with the measurement.

## 4 Experiments

Experimental setup.
For the CT experiments, we used the AAPM 2016 CT challenge data from McCollough et al. ().
We applied the same data processing methods as in Hu et al. ()
with the exception that
we used all the XY slices from the 9 training volumes
to train the in distribution networks, yielding a total of approximately 5000 slices.
For the deblurring and superresolution experiments,
we used the CelebA-HQ dataset (Liu et al., ) with each image having size $256\times 256$.
The test data was a randomly selected subset of 10 of the images not used for training.
In all cases, we report the average metrics across the test images:
peak SNR (PSNR) in dB, and structural similarity metric (SSIM) (Wang et al., ).
For the training data,
we trained networks on generated phantom images consisting of
randomly placed ellipses of different shapes and sizes.
See Fig. 20 for examples.
These phantoms can be generated on the fly in large quantities.
We used networks trained on grayscale phantoms for the CT experiments
and networks trained on RGB phantoms for the deblurring and superresolution experiments.
Appendix A.4 contains precise specifications of the phantoms.

We trained the patch-based networks with $64\times 64$ patches
and used a zero padding value of 64,
so that 5 patches in both directions were used to cover the target image.
We used the network architecture in Karras et al. ()
for both the patch-based networks and whole-image networks.
All networks were trained on PyTorch using the Adam optimizer with 2 A40 GPUs.

Single measurement setting.
In cases where no training data is available and we only have the measurement $\bm{y}$,
we applied Algorithm 1 to solve a variety of inverse problems:
CT reconstruction, deblurring, and superresolution.
For the forward and backward projectors in CT reconstruction,
we used the implementation provided by the ODL Team (). We performed two sparse-view CT (SVCT) experiments:
one using 20 projection views, and one using 60 projection views.
Both of these were done using a parallel beam forward projector
where the detector size was 512 pixels.
For the deblurring experiments,
we used a uniform blur kernel of size $9\times 9$
and added white Gaussian noise with $\sigma=0.01$
where the clean image was scaled between 0 and 1.
For the superresolution experiments, we used a scaling factor of 4
with downsampling by averaging
and added white Gaussian noise with $\sigma=0.01$.

For the comparison methods,
we ran experiments that naively used the OOD diffusion model
without the self-supervised network refining process.
For reference, we also ran experiments using a diffusion model
trained on the entire in-distribution training set (the “correct” model).
In practice, it would not be possible to obtain
such a large training dataset of in-distribution images.
Additionally, for these diffusion model methods,
we implemented both the patch-based version as well as the whole-image version.
The whole-image networks were trained with the loss function in (9)
and used the same network architecture as the patch-based models,
but the input of the network was the entire image
and did not contain positional encoding information.

We also compared with more traditional methods:
applying a simple baseline, reconstructing via the total variation regularizer (ADMM-TV),
and two plug and play (PnP) methods: PnP-ADMM (Xu et al., ) and PnP-RED (Hu et al., ).
For CT,
the baseline was obtained
by applying the filtered back-projection method to the measurement ${\bm{y}}$.
For deblurring,
the baseline was simply equal to the blurred image.
For superresolution,
the baseline was obtained by upsampling the low resolution image
and using nearest neighbor interpolation.
The implementation of ADMM-TV can be found in Hong et al. ().
Finally, since we assume we do not have access to any clean training data,
we used the off the shelf denoiser BM3D (Dabov et al., ).
Appendix A.5
contains the values of all the parameters of the algorithms.

Table 1 shows the main results
for single-measurement inverse problem solving.
The bottom two rows show the hypothetical performance
if it were possible to train a diffusion model on a large dataset of in distribution images,
which is not available in practice.
Our self-supervised patch-based diffusion approach
achieved significantly higher quantitative results
when averaged across the test dataset
than the self-supervised whole-image approach in all the inverse problems.
Furthermore, although the diffusion model that was initially used in this algorithm
was trained on completely different images, by applying the self-supervised loss,
the patch-based approach is able to achieve results that are close to
(and for the deblurring case, even surpassing) those using the in-distribution networks.
The table also shows that by including the self-supervised step,
a dramatic improvement over naively using the OOD model is achieved.
Lastly, Fig. 2 shows that some artifacts appear in the whole-image SS method
that are not present in our patch SS method.

Table: Table 1: Comparison of quantitative results on three different inverse problems in the single measurement setting. Results are averages across all images in the test dataset. Best results for practical use are in bold.

To demonstrate that our method also works well
even when the mismatched distribution is closer to the true distribution,
we also ran an experiment
where the networks were initially trained on the LIDC-IDRI dataset (Armato et al., ).
We extracted 10000 2D slices from the 3D volumes
and rescaled all the images so that the pixel values were between 0 and 1.
We then ran Algorithm 1 to perform CT reconstruction
where the test dataset was the same as the one used in Table 1.
Table 4 shows the results of this experiment.
Our method achieved better quantitative results than the whole image method
and even outperformed the reconstructions using the in distribution network
but without any self-supervision.
Appendix A.1 shows the visual results of these experiments.
Appendix A.2 further discusses using self-supervision
in cases where the initial network was trained on in-distribution data and shows improved image quality.

We also ran ablation studies to examine the effect of various parameters on the proposed method.
Barbano et al. () and Chung & Ye () used the LoRA module
for solving single-measurement inverse problems with diffusion models.
We tested this method for CT reconstruction and deblurring with different rank adjustments
and found this method to be inferior to modifying the weights of the entire network.
We also ran experiments using networks with different numbers of weights.
Appendix A.3
shows the results of these experiments.

Figure: Figure 2: Results of 60 view CT reconstruction using self supervised (SS) approach. The display uses modified HU units to show more contrast between organs.
Refer to caption: x2.png

Figure: Figure 3: Results of deblurring using self supervised (SS) approach and comparison methods.
Refer to caption: x3.png

Small dataset setting.
We ran experiments on the same inverse problems as the single measurement case.
The OOD networks were fine-tuned with 10 images randomly selected
from the in-distribution training set;
we also ran ablation studies using different quantities of in-distribution data
in Appendix A.3.
Figures 5 and 5 show that the patch-based model is much less prone to overfitting than the whole-image model.
Hence, to evaluate the best possible performance of the whole-image model compared to the patch-based model, for both models
we chose the checkpoint yielding the best results for solving inverse problems.

Table 2 shows the main results for solving inverse problems
using the fine-tuned diffusion model.
We compared the results of fine-tuning the whole-image model
with fine-tuning the patch-based model
as well as the best baseline out of the four baselines shown in Table 1.
The results show that the proposed patch-based method achieved the best performance
in terms of quantitative metrics for all of the inverse problems.
Figure 6 shows the visual results of this experiment.
The patch-based model is able to learn an acceptable prior
using the very small in-distribution dataset
and the reconstructed images contain fewer artifacts than the comparison methods.

Table: Table 2: Comparison of results for using diffusion models fine-tuned on 10 in-distribution images to solve inverse problems in small dataset setting. Best results are in bold.

Table 3
further investigates the effect of overfitting.
For different amounts of training time using the small in-distribution dataset,
we ran the reconstruction algorithm for 60-view CT.
While the whole-image model exhibited substantial image degradation
when the network was fine-tuned for too long,
the patch-based model retained relatively stable performance throughout the entire training process.
This illustrates that whole-image diffusion models exhibits severe overfitting problems
when only a small amount of training data is unavailable.
Furthermore, patch-based diffusion models assist greatly with this problem
and the results are evident for solving inverse problems.
Appendix A.1 shows the visual results of these experiments.

Figure: Figure 4: Comparison of PSNR between patch-based model and whole-image model for overfitting in small dataset setting.
Refer to caption: extracted/5929130/figs/overfitting.png

Figure: Figure 6: Results of inverse problem solving in the small dataset setting. Top row is 60 view CT recon, middle row is deblurring, and bottom row is superresolution. For CT, measurement refers to FBP.
Refer to caption: x4.png

To look at the priors learned by the different models from fine-tuning,
we unconditionally generated images from the checkpoints obtained by fine-tuning
on the 10 image CT dataset.
Figure 7 shows a subset of the generated images
where we used the checkpoints obtained after 4 hours of training.
The top two rows consist of images generated by the whole-image model
and the bottom two rows consist of images generated by the patch diffusion model.
To emphasize the memorization point,
we grouped together similar looking images in the top two rows:
it can be seen that the images in each group look virtually identical,
despite the fact that the pure white noise initializations for each sample was different.
On the other hand, while the samples generated by the patch diffusion model
also show some unrealistic features,
they all show some distinct features,
which implies that this model has much better generalization ability.

Figure: Figure 7: Unconditional generation of CT images from networks fine-tuned in the small dataset setting. Top two rows were generated with the whole image model; bottom two rows were generated with the patch-based model.
Refer to caption: x5.png

## 5 Conclusion

This paper presented a method of using patch-based diffusion models
to solve inverse problems when the data distribution is mismatched from the trained network.
In particular, we conducted experiments in the setting when only a single measurement is available as well as the setting when a very small subset of in-distribution data is available.
In both settings, the proposed patch-based method outperformd whole-image methods
in a variety of inverse problems.
In the future, more work could be done on using acceleration methods for faster reconstruction,
exploring other less computationally expensive methods of fine-tuning the network
geared toward inverse problem solving,
and methods of refining the prior when a set of measurements are available
(Yaman et al., ).
Limitations of the work include a slow runtime for the self-supervised algorithm
and a lack of theoretical guarantees for convergence of algorithms and dataset size requirements.

### Acknowledgments

## Appendix A Appendix

### A.1 Additional inverse problem solving figures

Figure 8 shows the results of various methods applied to superresolution in the single measurement setting.

Figure: Figure 8: Results of superresolution using self supervised (SS) approach and comparison methods.
Refer to caption: x6.png

Figure 9 shows the results of 20 view CT reconstruction using Algorithm 1.
This very sparse view CT recon problem is made more challenging by the lack of any training data.
Artifacts can clearly be seen in all the comparison methods.
Despite this challenge, reconstructions such as this one can still be useful for medical applications such as patient positioning.

Figure: Figure 9: Results of 20 view CT reconstruction in a self-supervised setting. For clarity, the images are plotted on the same scale as the diffusion models were trained.
Refer to caption: x7.png

Figure 10 shows the results of running self-supervised CT reconstruction with 20 views and 60 views
where the starting checkpoint was obtained through training on a large (but out of distribution) CT dataset:
10000 LIDC-IDRI slices.
Particularly for 20 views, the artifacts from using the whole image model are apparent,
while the patch-based model obtains a much higher quality reconstruction.
Thus, regardless of whether the starting network has a severely mismatched distribution (ellipses)
or a slightly mismatched distribution (different CT dataset), our proposed method outperforms the whole image model.

Figure: Figure 10: Results of CT reconstruction in a self-supervised setting when the starting network was trained on the LIDC dataset. Top row used 20 views and bottom row used 60 views.
Refer to caption: x8.png

Figure 11 shows the results of performing 60 view CT reconstruction
in an unsupervised manner
from checkpoints fine-tuned using the small in distribution CT dataset.
The images on the bottom row shows the progressively worsening degradation
and increasing number of artifacts
resulting from overfitting exhibited by whole image model.
On the other hand, the top row shows relatively stable performance exhibited by the patch-based model
as it is able to avoid overfitting much better.

**Table 3: Performance of fine-tuning on 60 view CT using checkpoints trained for different lengths of time. Best results are in bold.**
| Train | Patches | Whole image |  |  |
| --- | --- | --- | --- | --- |
| time (hr) | PSNR$\uparrow$ | SSIM $\uparrow$ | PSNR$\uparrow$ | SSIM $\uparrow$ |
| 0 | 33.91 | 0.921 | 33.10 | 0.911 |
| 0.5 | 40.37 | 0.964 | 38.39 | 0.959 |
| 1 | 40.91 | 0.965 | 40.54 | 0.964 |
| 2 | 41.21 | 0.965 | 39.19 | 0.953 |
| 3 | 41.17 | 0.965 | 38.31 | 0.945 |
| 4 | 41.02 | 0.964 | 37.67 | 0.938 |

Figure: Figure 11: Results of 60 view CT recon with networks fine-tuned on 10 in distribution CT images for varying amounts of training time.
Refer to caption: x9.png

### A.2 Effect of self-supervision for different distributions

Recall that in the single measurement setting,
Algorithm 1 is used to adjust the underlying distribution of the network
away from the originally trained OOD data and toward the ground truth image.
We investigated the effect of applying this method
even when the network was trained on the in-distribution data.
Figures 13 and 13
show the results of this experiment for CT reconstruction,
where each point represents the specific PSNR for one of the images in the test dataset.
If the additional self-supervision step had no effect on the image quality,
the points would lie on the red line.
However, in both cases, all of the points are above the red line,
indicating that the self-supervision step of the algorithm
improves the image quality even when the network was already trained on in-distribution data.
Furthermore, the improvement is more substantial for for the 20 view case than the 60 view case,
as the predicted clean images $D_{\theta}({\bm{x}}_{t}|{\bm{y}})$ at each step for the 60 view case
are likely to be more closely aligned with the measurement,
so the network refining step becomes less significant.
Importantly, this shows that in practice,
one may directly apply Algorithm 1 to solve inverse problems
without knowledge of the severity of the mismatch in distribution between training and testing data:
even when there is no mismatch, the additional self-supervision step does not degrade the image quality.

Figure: Figure 12: PSNR of 20 view CT reconstruction in single-measurement setting using a patch-based in-distribution network.
Refer to caption: extracted/5929130/figs/Figure_1_20view.png

Table 5 summarizes the results of using different training datasets
while keeping the same test dataset (AAPM CT images).
The distribution shift is greatest when the network is trained on ellipse phantoms
and used to reconstruct the AAPM CT images, so the reconstruction quality is the lowest in this case.
The LIDC dataset consists of CT images which belong to a distribution
that is reasonably similar to the distribution of AAPM CT images,
so when using the network trained on LIDC images,
the quality drop over using an in-distribution network is not substantial.
Finally, the improvements obtained by using more in-distribution networks
is more apparent for the 20 view case as the measurements are sparser for this case,
so the prior plays a larger role in obtaining an accurate reconstruction.

**Table 4: Single measurement CT reconstruction results where the initial checkpoint was trained on LIDC dataset and refined on the fly with the AAPM measurement.**
| Dataset | CT, 20 views | CT, 60 views |  |  |
| --- | --- | --- | --- | --- |
| size | PSNR$\uparrow$ | SSIM $\uparrow$ | PSNR$\uparrow$ | SSIM $\uparrow$ |
| Whole image | 35.01 | 0.894 | 41.95 | 0.967 |
| Patches (Ours) | 36.34 | 0.918 | 42.32 | 0.972 |

**Table 5: Performance of patch-based model in single measurement setting for CT reconstruction for different OOD training datasets.**
| Train | 20 views | 60 views |  |  |
| --- | --- | --- | --- | --- |
| time (hr) | PSNR$\uparrow$ | SSIM $\uparrow$ | PSNR$\uparrow$ | SSIM $\uparrow$ |
| Ellipses | 33.77 | 0.874 | 41.45 | 0.966 |
| LIDC | 36.34 | 0.918 | 42.32 | 0.970 |
| AAPM | 36.82 | 0.923 | 42.33 | 0.970 |

### A.3 Ablation studies

We performed four ablation studies
to evaluate the impact of various parameters on the proposed methods.
Similar to the main text,
all quantitative results are averaged across the test dataset.

Low rank adaptation.
To avoid overfitting to the measurement in self-supervised settings,
Barbano et al. () proposed using a low rank adaptation to the weights of the neural network,
reducing the number of weights that are adjusted during reconstruction by a factor of around 100.
Here we investigate the effect of using different ranks of adaptations on two inverse problems:
60 view CT reconstruction and deblurring.
Consistent with Barbano et al. () and Chung & Ye (),
we only used the LoRA module for attention and convolution layers.
We also allowed the biases of the network to be changed.

Tables 6 and 7 show the quantitative results of these experiments,
where a rank of “full” represents fine-tuning all the weights of the network.
In all cases, using LoRA for this fine-tuning process results in worse reconstructions than simply fine-tuning the entire network.
The visual results are especially apparent in Figure 15:
the reconstructed image becomes oversmoothed when using LoRA and artifacts become present when using the whole image model.
This is likely due to the large distribution shift between the initial distribution of images and target distribution of faces:
the low rank adaptation to the mismatched network is not sufficient to represent the new distribution
and thus the self-supervised loss function results in smoothed images.

**Table 6: Performance of 60 view CT recon using self-supervised network refining with LoRA module. Best results are in bold.**
| Rank | Parameters (%) | Patches | Whole image |  |  |
| --- | --- | --- | --- | --- | --- |
|  |  | PSNR$\uparrow$ | SSIM $\uparrow$ | PSNR$\uparrow$ | SSIM $\uparrow$ |
| 2 | 1.1 | 40.37 | 0.963 | 39.25 | 0.952 |
| 4 | 2.0 | 40.32 | 0.963 | 39.10 | 0.951 |
| 8 | 3.8 | 40.33 | 0.963 | 39.18 | 0.951 |
| 16 | 7.2 | 40.32 | 0.963 | 39.33 | 0.953 |
| Full | 100 | 41.45 | 0.966 | 40.47 | 0.957 |

**Table 7: Performance of deblurring using self-supervised network refining with LoRA module. Best results are in bold.**
| Rank | Parameters (%) | Patches | Whole image |  |  |
| --- | --- | --- | --- | --- | --- |
|  |  | PSNR$\uparrow$ | SSIM $\uparrow$ | PSNR$\uparrow$ | SSIM $\uparrow$ |
| 2 | 1.1 | 29.31 | 0.830 | 29.19 | 0.811 |
| 4 | 2.0 | 29.31 | 0.829 | 29.35 | 0.817 |
| 8 | 3.8 | 29.38 | 0.831 | 29.19 | 0.810 |
| 16 | 7.2 | 29.31 | 0.830 | 29.33 | 0.815 |
| Full | 100 | 30.34 | 0.860 | 29.50 | 0.831 |

Figure: Figure 14: Results of using LoRA module for 60 view CT reconstruction in a single measurement setting. All weights refers to adjusting all the weights of the network at reconstruction time.
Refer to caption: x10.png

Figure: Figure 15: Results of using LoRA module for deblurring in a single measurement setting. All weights refers to adjusting all the weights of the network at reconstruction time.
Refer to caption: x11.png

Effect of network size.
In the self-supervised case, another potential method to avoid overfitting
is to use a smaller network.
We trained networks with differing numbers of base channels (but no other modifications)
on the ellipse phantom dataset
and then used Algorithm 1 to perform self-supervised 60 view CT reconstruction.
Table 8 shows the quantitative results of this experiment.
For both the patch-based model and the whole image model,
the network with 128 base channels obtained the best result,
so we used this network architecture for all the main experiments.
Figure 16again shows evidence of overfitting in the form of artifacts
in the otherwise smooth regions of the organs when using the network with 256 base channels.
These artifacts are less obvious in the patch-based model.

**Table 8: Performance of 60 view CT recon in a self-supervised manner with networks of different sizes. Best results are in bold.**
| Base | Parameters | Patches | Whole image |  |  |
| --- | --- | --- | --- | --- | --- |
| Channels | (Millions) | PSNR$\uparrow$ | SSIM $\uparrow$ | PSNR$\uparrow$ | SSIM $\uparrow$ |
| 32 | 3.4 | 39.73 | 0.958 | 39.69 | 0.957 |
| 64 | 14 | 40.37 | 0.961 | 40.07 | 0.958 |
| 128 | 60 | 41.45 | 0.966 | 40.47 | 0.957 |
| 256 | 217 | 40.29 | 0.959 | 39.28 | 0.954 |

Figure: Figure 16: Results of 60 view CT recon using networks with different numbers of parameters in the single-measurement setting. The top numbers show the number of total parameters in the network.
Refer to caption: x12.png

Fine-tuning with a larger dataset.
To examine the effect of fine-tuning the networks on differing sizes of in-distribution datasets,
we started with the same checkpoint trained on ellipses
and fine-tuned them using various sizes of datasets consisting of CT images.
Each small dataset consisted of randomly selected images from the entire 5000 image AAPM dataset.
Next we used these checkpoints to perform 60 view CT reconstruction (without any self supervision).
Table 9 shows the results of these experiments,
where we also included the results of using the in-distribution network
trained on the entire 5000 image dataset.
This shows that for a wide range of different fine-tuning dataset sizes
our proposed method obtained better metrics than the whole-image model.

**Table 9: Performance of fine-tuning on 60 view CT using checkpoints fine-tuned from different dataset sizes. Best results are in bold.**
| Dataset | Patches | Whole image |  |  |
| --- | --- | --- | --- | --- |
| size | PSNR$\uparrow$ | SSIM $\uparrow$ | PSNR$\uparrow$ | SSIM $\uparrow$ |
| 3 | 40.93 | 0.964 | 40.45 | 0.964 |
| 10 | 41.21 | 0.965 | 40.54 | 0.964 |
| 30 | 41.31 | 0.966 | 40.66 | 0.967 |
| 100 | 41.46 | 0.967 | 40.96 | 0.968 |
| 5000* | 41.70 | 0.967 | 41.67 | 0.969 |

Figure: Figure 17: Results of 60 view CT recon in the small dataset setting where the size of the small dataset is varied.
Refer to caption: x13.png

Backpropagation iterations during self-supervision.
In the single measurement setting, the self-supervised loss is crucial
to ensuring that the OOD network output is consistent with the measurement.
Backpropagation through the network is necessary to minimize this loss,
but too much network refining during this step could lead to overfitting
to the measurement and image degradation.
We ran experiments examining the effect of the number of backpropagation iterations
during each step for the patch-based model and the whole image model.
Figures 19 and 19 show that in both cases,
performance generally improved when increasing the number of backpropagation iterations
and overfitting is avoided.
Additionally, the patch-based model always outperformed the whole image model
and exhibited more improvement as the number of backpropagation iterations increased.
For our main experiments,
we used 5 iterations as the improved performance became marginal compared to the extra runtime.

Figure: Figure 18: Comparison of PSNR between patch-based model and whole-image model for number of network refining iterations in single measurement setting.
Refer to caption: extracted/5929130/figs/ablation7_psnr.png

**Table 10: Performance of Algorithm 1 for 60 view CT reconstruction in single measurement setting with different numbers of backpropagation iterations. Best results are in bold.**
| Backprop | Patches | Whole image |  |  |
| --- | --- | --- | --- | --- |
| iterations | PSNR$\uparrow$ | SSIM $\uparrow$ | PSNR$\uparrow$ | SSIM $\uparrow$ |
| 0 | 33.97 | 0.934 | 33.10 | 0.911 |
| 1 | 40.35 | 0.964 | 39.81 | 0.958 |
| 2 | 40.96 | 0.966 | 40.45 | 0.961 |
| 5 | 41.45 | 0.966 | 40.47 | 0.957 |
| 10 | 41.65 | 0.968 | 40.54 | 0.958 |
| 20 | 41.92 | 0.970 | 40.71 | 0.959 |
| 50 | 42.18 | 0.971 | 40.90 | 0.961 |

### A.4 Phantom dataset details

We used two phantom datasets of 10000 images each:
one consisting of grayscale phantoms and the other consisting of colored phantoms.
The grayscale phantoms consisted of 20 ellipses with a random center within the image,
each with minor and major axis having length equal to a random number
chosen between 2 and 20 percent of the width of the image.
The grayscale value of each ellipse was randomly chosen between 0.1 and 0.5;
if two or more ellipses overlapped,
the grayscale values were summed for the overlapped area with all values exceeding 1 set to 1.
Finally, all ellipses were set to a random angle of rotation.
The colored phantoms were generated in the same way,
except the RGB values for each ellipse were set independently
and then multiplied by 255 at the end.
Figure 20 shows some of the sample phantoms.

Figure: (a) Six grayscale phantoms
Refer to caption: x14.png

### A.5 Experiment parameters

We applied the framework of Karras et al. () to train the patch-based networks and whole image networks.
Since images were scaled between 0 and 1 for both grayscale images and RGB channels,
we chose a maximum noise level of $\sigma=40$
and a minimum noise level of $\sigma=0.002$ for training.
We used the same UNet architecture for all the networks
consisting of a base channel multiplier size of 128
and 2, 2, and 2 channels per resolution for the three layers.
We also used dropout connections with a probability of 0.05
and exponential moving average for weight decay
with a half life of 500K images to avoid overfitting.

The learning rate was chosen to be $2\cdot 10^{-4}$ when training networks from scratch
and was $1\cdot 10^{-4}$ for the fine-tuning experiments.
For the patch-based networks, the batch size for the main patch size ($64\times 64$) was 128,
although batch sizes of 256 and 512 were used for the two smaller patch sizes of $32\times 32$ and $16\times 16$. The probabilities of using these three patch sizes were $0.5,0.3$, and $0.2$ respectively.
For the whole image model, we kept all the parameters the same,
but used a batch size of 8.

For image generation and inverse problem solving,
we used a geometrically spaced descending noise level
that was fine tuned to optimize the performance for each type of problem.
We used the same set of parameters
for the patch-based model and whole image model.
The values without the self-supervised loss are as follows:

- •
CT with 20 and 60 views: $\sigma_{\max}=10,\sigma_{\min}=0.005$
- •
Deblurring: $\sigma_{\max}=40,\sigma_{\min}=0.005$
- •
Superresolution: $\sigma_{\max}=40,\sigma_{\min}=0.01$.

The values with the self-supervised loss are as follows:

- •
CT with 20 and 60 views: $\sigma_{\max}=10,\sigma_{\min}=0.01$
- •
Deblurring: $\sigma_{\max}=1,\sigma_{\min}=0.01$
- •
Superresolution: $\sigma_{\max}=1,\sigma_{\min}=0.01$.

Finally, for generating the CT images we used $\sigma_{\max}=40,\sigma_{\min}=0.005$.

When running Algorithm 1,
we set $K=10$ for all experiments and $M=5$ for CT reconstruction
and $M=1$ for deblurring and superresolution.
We ran 5 iterations of network backpropagation with a learning rate of $10^{-5}$.
When using the LoRA module as in the ablation studies (see Tables 7 and 6),
we ran 10 iterations of network backpropagation with a learning rate of $10^{-3}$.

The ADMM-TV method for linear inverse problems consists of solving the optimization problem

$$ $\text{argmax}_{{\bm{x}}}\frac{1}{2}\|{\bm{y}}-A{\bm{x}}\|_{2}^{2}+\lambda\, \text{TV}({\bm{x}}),$ (13) $$

where $\text{TV}({\bm{x}})$ represents the L1 norm total variation of $vx$,
and the problem is solved with the alternating direction method of multipliers.
For CT reconstruction, deblurring, and superresolution,
we chose $\lambda$ to be $0.001,0.002$, and $0.006$ respectively.

The PnP-ADMM method consists of solving the intermediate optimization problem

$$ $\text{argmax}_{{\bm{x}}}f({\bm{x}})+(\rho/2)\|{\bm{x}}-({\bm{z}}-{\bm{u}})\|_{ 2}^{2},$ (14) $$

where $\rho$ is a constant.
The values for $\rho$ we used for CT reconstruction, deblurring, and superresolution
were 0.05, 0.1, and 0.1 respectively.
We used BM3D as the denoiser with a parameter representing the noise level:
this parameter was set to 0.02 for 60 view CT and 0.05 for the other inverse problems.
A maximum of 50 iterations of conjugate gradient descent was run per outer loop.
The entire algorithm was run for 100 outer iterations at maximum
and the PSNR was observed to decrease by less than 0.005dB per iteration by the end.

The PnP-RED method consists of the update step

$$ ${\bm{x}}\leftarrow{\bm{x}}+\mu(\nabla f-\lambda({\bm{x}}-D({\bm{x}}))),$ (15) $$

where $D({\bm{x}})$ represents a denoiser.
The stepsize $\mu$ was set to 0.01 for the CT experiments and 1 for deblurring and superresolution.
We set $\lambda$ to 0.01 for the CT experiments and 0.2 for deblurring and superresolution.
Finally, the denoiser was kept the same as the PnP-ADMM experiments with the same denoising strength.

Table 11 shows the average runtimes of each of the implemented methods
when averaged across the test dataset for 60 view CT reconstruction.

**Table 11: Average runtimes of different methods across images in the test dataset for 60 view CT recon.**
| Method | Runtime (s) $\downarrow$ |
| --- | --- |
| Baseline | 0.1 |
| ADMM-TV | 1 |
| PnP-ADMM | 73 |
| PnP-RED | 121 |
| Whole diffusion | 112 |
| Whole SS | 248 |
| Whole LoRA | 329 |
| Patch diffusion | 123 |
| Patch SS | 289 |
| Patch LoRA | 377 |

### A.6 Self-supervised Inverse Problem Figures

The following figures show additional examples of self-supervised inverse problem solving.

Figure 21 shows additional example slices of CT reconstruction from 60 views.

Figure 22 shows additional example slices of CT reconstruction from 20 views.

Figure 23 shows additional examples of deblurring with face images.

Figure 24 shows additional examples of superresolution with face images.

Figure: Figure 21: Additional figures for self-supervised 60 view CT recon.
Refer to caption: x16.png

Figure: Figure 22: Additional figures for self-supervised 20 view CT recon.
Refer to caption: x17.png

Figure: Figure 23: Additional figures for self-supervised deblurring.
Refer to caption: x18.png

Figure: Figure 24: Additional figures for self-supervised superresolution.
Refer to caption: x19.png

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: - In their appendix, [1] provide experimental results when training whole image and patch-based models on training datasets of different sizes, and similarly observe that PaDIS remains visually more consistent in contrast to the whole-image models. Could the authors elaborate how their experiments and observations relate and potentially complement those of [1]?
- See also the weaknesses above.
soundness: 3
strengths: - The considered domain adaption tasks are important for practical application.
- The observation that patch-based diffusion models are more robust to finetuning / out-of-distribution tasks is interesting and relevant.
summary: The paper considers the problem of adapting diffusion models trained on a domain A  to the task of solving reconstruction problems on another domain B, and investigate the cases where a single measurement from B or a small number of samples from B are available.
To that end, the authors investigate deep diffusion image prior (DDIP) adaption methods, which were originally proposed to be used with whole-image diffusion models, and combine with the recently proposed patch-based diffusion models. The experimental results on CT and natural image datasets shows that patch-based diffusion models are more robust to fine-tuning on small dataset as well as to adapting the network weights for single-measurement domain adaption.
weaknesses: - While it is important to improve upon the considered tasks, and the paper presents good results on that task, I am unsure about the contribution of the paper. The paper combines the existing SCD/DDIP method with the recently proposed patch-based diffusion models and their inverse-problem solver (PaDIS), but it seems that this is only a matter of replacing the whole-image with patch-based diffusion models and their score calculation method. The robustness to OOD data of patch-based models is important, but the relation to the experiments on small datasets presented by [1] is unclear.
- The results in Table 1 show that the main performance gains are due to using SCD/DDIP and using that with Patch-based DMs increases the PSNR by at most 1dB. While this is certainly an improvement, using the patch-based model in-distribution already increases the PSNR.
- In section 3, the authors introduce the patch-based diffusion prior method, where they seem to have copied Figure 1 and the text (with small adaptions) from the original paper. I think this should be stated more clearly. Moreover, some parts in the argumentation seem to missing. As an example, L.198-199: "represents the aforementioned bordering region", but the bordering region has not really been mentioned before (in contrast to the original text). Restating the assumed probability distribution in L.195, but it would be helpful to also state the motivation (the calculation of the score model based on the patch scores).

References:
- [1]: Hu et al (2024): Learning Image Priors through Patch-based Diffusion Models for Solving Inverse Problems

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
questions: I would like to see what the results would be like when applying the method to the black whole imaging problem [1] where the true prior is unavailable.

[1] Wu, Zihui, et al. "Principled Probabilistic Imaging using Diffusion Models as Plug-and-Play Priors." arXiv preprint arXiv:2405.18782 (2024).
soundness: 3
strengths: 1. The writing is clear so easy to follow.

2. The motivation of using patch-based prior for better generalizability is reasonable.

3. The proposed method addresses an important practical problem.
summary: The paper proposes to use patch-based diffusion models for solving inverse problems with mismatched training and test distributions. The authors address the challenge of artifacts and hallucinations in image reconstructions when the training and test datasets are not aligned. They propose a patch-based approach that leverages image patches to learn priors, demonstrating the effectiveness of the proposed method in scenarios with limited data availability.
weaknesses: 1. It is not clear why Eq. (11) can address "The image that is being reconstructed might not come from the distribution of the training images". I recommend the authors to provide more detailed discussion.

2. The proposed method is similar to Deep Diffusion Image Prior for Efficient OOD Adaptation in 3D Inverse Problems. It could be beneficial to discuss the similarity and difference to highlight the contribution of this paper.

---

confidence: 5
contribution: 2
initial_rating: 3
presentation: 2
questions: 1. Can the authors provide some figures of training data? I have seen Fig. 20 but I still can not imagine the training data.
2. Why Fig. 6 does not have the same notations as Fig. 3? It looks really confusing.
3. I am wondering if the proposed method can scale up. Since patch-based diffusion models should aim for large images, the experiments on CelebA 256*256 are insufficient.

I am willing to raise the score if the concerns are well-addressed.
soundness: 3
strengths: 1. The paper is easy to follow.
2. The topic of mismatching distribution inverse problems is timely and important.
3. Using patch-based diffusion models for mismatching distribution inverse problems is natural.
summary: This paper shows that the patch-based diffusion model can be a good solution for mismatching distribution inverse problems, compared to the conventional whole-image diffusion model. The authors study this setting where (1) only measurements are available and (2) very small ID samples are available, the results are decent.
weaknesses: 1. I doubt the contribution of the work. As it is not new to solve inverse problems with patch-based diffusion models[1], the finding of this paper 'whole-image models are prone to memorization and overfitting, while a patch-based model can resolve these issues' is already clarified in [1].

2. Lack of theoretical analysis. It is straightforward that patch-based diffusion models are suited for mismatching distribution inverse problems, as they can avoid memorization and overfitting ID data. It would be good if the authors could provide some theory on this argument.

3. Minor: 

      Errors in Eq. 2 and Eq. 3; 

      Better to repaint Fig. 1 instead of directly copying it from [1] without citation.


[1] Learning Image Priors through Patch-based Diffusion Models for Solving Inverse Problems. Jason Hu, et al.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 1
questions: Please provide clarity on points raised under the weaknesses section.
soundness: 2
strengths: - The paper addresses an important and practical problem that commonly arises in real-world scenarios, where the generative prior distribution differs from the target distribution. This focus on out-of-distribution (OOD) issues has strong applicability in diverse settings.

- The experimental results show that the proposed patch-based model achieves better performance compared to whole-image models, underscoring the effectiveness of a patch-based approach in handling OOD challenges for inverse problems.
summary: This paper examines the use of diffusion models in inverse problems, particularly when there’s a mismatch between the training and test data distributions (out-of-distribution, OOD). The authors investigate two settings: one where only a single measurement from an unknown test distribution is available, and another where a small sample from the test distribution is accessible. They propose a patch-based diffusion model that learns image priors from patches rather than entire images. This model includes a self-supervised loss to enhance consistency with the given measurement in the single-measurement scenario. Their experiments demonstrate that this patch-based approach produces high-quality reconstructions.
weaknesses: - Limited novelty. I believe that the paper’s primary contribution—a combination of a patch-based model with self-supervised loss (e.g., deep image priors) to address out-of-distribution (OOD) issues in inverse problems—builds on existing concepts. While the integration of these components to tackle a specific challenge is interesting, the novelty is somewhat limited, as each component’s effectiveness has been demonstrated in prior works. Additionally, there is a lack of theoretical justification or clear intuition for the design of Algorithm 1, which would strengthen the proposed approach.

- The experimental setup lacks clarity. Specifically, it is not explained how the authors achieve the “Whole image, correct*” model or which sampling algorithm is used. Key hypothetical models are described ambiguously, and the term “best baselines” in Table 2 is undefined, making it challenging to understand the comparisons being drawn.

- Lack of analysis. Why the proposed algorithm 1 or patch-based model is better than previous models is not clearly demonstrated in the paper.

- Missing baselines. The paper does not include comparisons with strong baseline methods for inverse problem-solving using diffusion models, such as DPS [1], DDS [2], DDNM [3], and DAVI [4]. These methods are relevant in both single-measurement and limited-sample settings. In particular, [4] addresses OOD settings, making it a highly relevant benchmark for this work.

[1] Diffusion Posterior Sampling for General Noisy Inverse Problems, ICRL23 \
[2] Decomposed Diffusion Sampler for Accelerating Large-Scale Inverse Problems, ICLR24 \
[3] Zero-Shot Image Restoration Using Denoising Diffusion Null-Space Model, ICLR23 \
[4] Diffusion Prior-Based Amortized Variational Inference for Noisy Inverse Problems, ECCV24

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 1
questions: Questions/feedback on clarity: 

- 3.1. is difficult to follow without already knowing the framework the authors adapt. Why is the bordering region added? What is $M$ here, and what is $k$? What does $i$ denote? Is $x$ an image patch or the whole image, as it has been used throughout the paper for both. Why only the x-positions of patches are concatenated as input, why not both x and y coordinates?

- Probably due to the unclear nature of 3.1. I was unable to properly follow 3.2 in some parts. What do authors mean by "the outermost product is computationally very expensive"?

Questions/feedback on experiments:

- How many samples have been used to produce the results reported in Table 1?

- I recommend reporting perceptual metrics such as LPIPS as well, especially for image deblurring and superresolution.

- How is it possible that the proposed method, without training, outperforms diffusion approaches that leverage training data? This sounds very counter-intuitive.

- Which dataset has been used to produce Figures 4 and 5?

- The discussion on diversity of generated samples is very vague. What do authors mean by samples "show some unrealistic features"? IT is unclear based on the Figure 7 which features are considered realistic/unrealistic. Claims about sample diversity would be more convincing if authors reported specific metrics about diversity, such as Recall.

- Table 3 is in the appendix, and therefore it should either moved to the main paper or the discussion about Table 3 should be moved to the appendix.
soundness: 2
strengths: - The problem is well motivated and important. In a multitude of real-world applications, it is infeasible to collect large datasets for training diffusion models on the target distribution. However, finding ways to leverage the strong prior provided by diffusion models to tackle such problems is a valuable direction.

- The key idea of using patch-based diffusion instead of whole-image diffusion is intuitive and sensible. The patch-based model is less prone to overfitting to the available source data distribution, and thus can perform better in the presence of distribution shifts. Moreover, patch-based models can be trained more data-efficiently which is crucial in the data scarce domain.

- The experimental results, if verified on larger scale experiments, are promising.
summary: This works tackles the challenging problem of solving inverse problems where only a few samples are available from the test distribution. Authors propose a patch-based diffusion prior in this scenario, that learns the image distribution from patches, and not from whole images. Authors argue that whole-image models are prone to overfitting to the data distribution, and thus are unable to provide sufficient performance when the test samples come from out of distribution. Numerical experiments are provided to verify these findings.
weaknesses: - The clarity of the paper could be greatly improved. In particular, I had a difficult time following 3.1, which is central to understanding the proposed algorithm and led to downstream confusions throughout the paper. Specific questions follow under 'Questions'.

- If I understand correctly, the experimental results are reported on 10 samples. This is not enough to report statistically meaningful results.

- Many claims in the experimental section are vague or not supported properly by the experiments (see more details under 'Questions').