Equivariant plug-and-play image reconstruction

August 25, 2025 · View on GitHub

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

To appear in CVPR 2024

tl;dr

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

flexible

Method description

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)

x_{k+1} = \text{D}(x_k - \gamma \nabla f(x_k)),

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

\begin{align*} &\text{Sample } g_k \sim \mathcal{G} \\ &\text{Set } \text{D}_{g_k}(x) = T_{g_k}^{-1} \text{D}(T_{g_k} x) \\ &x_{k+1} = \text{D}_{g_k}(x_k - \gamma \nabla f(x_k)). \end{align*}

Code

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:

PnP (click to expand)

without equivariance

On the set3c dataset for the motion blur problem, with the drunet model:


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

with equivariance

On the set3c dataset for the motion blur problem, with the drunet model, and with equivariance:


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

RED (click to expand)

without equivariance

On the set3c dataset for the super-resolution blur problem, with the drunet model (Fig. 6 of the paper):

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

with equivariance

On the set3c dataset for the motion blur problem, with the drunet model, and with equivariance (Fig. 6 of the paper):

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

Feel free to change problem and models!

ULA (click to expand)

without equivariance

On the BSD68 dataset for the super-resolution blur problem, with the drunet model (Fig. 8 of the paper):

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

with equivariance

On the BSD10 dataset for the motion blur problem, with the drunet model, and with equivariance (Fig. 8 of the paper):

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

Feel free to change problem and models!

Requirements

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.