camera-ready (soon!)

Matthieu Terris, Thomas Moreau, Nelly Pustelnik, Julián Tachella.

To appear in CVPR 2024

Enforcing equivariance of the denoiser to certain transformations within PnP/RED algorithms improves the stability and reconstruction quality of the algorithm.

We consider algorithms where gradients (or proximal operators) of explicit priors are replaced by denoisers; these algorithms typically take the form (in the case of PnP)

where $\text{D}$ is a denoiser. In our paper, we show that enforcing equivariance of the denoiser with respect to a group of geometric transformations (such as rotations) can improve the Lipschitz constant of the denoiser, and hence the stability of the algorithm. While a trivial way to enforce equivariance with respect to a group of transforms $\mathcal{G}$ is to perform an averaging of the denoiser's output over the group of transformations as $\text{D}_{\mathcal{G}} = \frac{1}{|G|} \sum_{g \in \mathcal{G}} T_g^{-1} \text{D}(T_g)$, we propose to use a Monte-Carlo estimation of the equivariant denoiser at each step of the algorithm. The resulting algorithm reads

To reproduce the experiments, first download the test datasets and place them in your data folder. Next, update the config/config.json file to point to the correct data folder. There, there are two folders to specify:

• ROOT_DATASET: the folder within which the CBSD68 and set3c datasets are located;
• PATH_MRI_DATA: the path to the fastMRI .pt dataset.

Then, you can run the following scripts to reproduce the experiments:

python running_pnp.py --problem='motion_blur' --model_name='drunet' --rand_rotations=0 --dataset_name='set3c' --results_folder='table_4/' --compute_lip=0 --sigma_den=0.02 --noise_level=0.01

python running_pnp.py --problem='motion_blur' --model_name='drunet' --rand_rotations=1 --dataset_name='set3c' --results_folder='table_4/' --compute_lip=0 --sigma_den=0.02 --noise_level=0.01

python running_red.py --problem='sr' --model_name='drunet' --rand_translations=0 --dataset_name='set3c' --sigma_den=0.015 --sr=2

python running_red.py --problem='sr' --model_name='drunet' --rand_translations=1 --dataset_name='set3c' --sigma_den=0.015 --sr=2

python running_ula.py --problem='motion_blur' --model_name='drunet' --rand_translations=0 --dataset_name='subset_BSD20' --sigma_den=0.019

python running_ula.py --problem='motion_blur' --model_name='drunet' --rand_translations=1 --dataset_name='subset_BSD20' --sigma_den=0.019

This code was tested with the following packages:

• torch 2.2
• deepinverse 0.1.1

The deepinverse package can be installed with pip install deepinverse or by cloning the repository.