UNSURE: Unknown Noise level Stein's Unbiased Risk Estimator

Code of the paper "UNSURE: Unknown Noise level Stein's Unbiased Risk Estimator" by Julian Tachella, Mike Davies and Laurent Jacques.

We use the deepinv library for most of the code.

The UNSURE loss was added to the deepinv library, please see this jupyter notebook demo and the documentation.

Paper available at arXiv.

Method Description

UNSURE is a self-supervised learning loss that can be used for learning a reconstruction network $f$ from a dataset of noisy measurements

$$ y_i = \mathcal{S}(Ax_i) $$

for $i=1,\dots,N$ where $x_i$ is the clean image, $y_i$ is the noisy measurement, $A$ is a linear operator and $\mathcal{S}$ is a stochastic noising process.

Unlike Stein's Unbiased Risk Estimator (SURE), the proposed loss can be used without any prior knowledge of the noise level. The loss is defined as

$$ \max_{\eta} \min_{f} \sum_{i=1}^N ||y_i - f(y_i)||^2 + 2\Sigma_{\eta} \text{div} f(y_i) $$

where $\eta$ is a Lagrange multiplier, $\Sigma_{\eta}$ is the covariance matrix of the noise and $\text{div} f(y)$ is the divergence of the network $f$ at the point $y$. There exist multiple variants for different types of noise (e.g. Gaussian, Poisson-Gaussian, spatially correlated noise, etc.). Please refer to the paper for more details.

Getting Started

1. Clone the repository
2. Install the latest version of deepinv if you don't have it already

pip install git+https://github.com/deepinv/deepinv.git#egg=deepinv

3. Generate the datasets by running the generate_datasets.py file.
4. Run the main.py file.

Citation

@misc{tachella2024unsureunknownnoiselevel,
      title={UNSURE: Unknown Noise level Stein's Unbiased Risk Estimator}, 
      author={Julián Tachella and Mike Davies and Laurent Jacques},
      year={2024},
      eprint={2409.01985},
      archivePrefix={arXiv},
      primaryClass={stat.ML},
      url={https://arxiv.org/abs/2409.01985}, 
}