Pre-Trained Denoisers for Inverse Problems

Definition:

Plug-and-Play (PnP) Framework

PnP replaces the proximal operator in iterative algorithms with a pre-trained denoiser DσD_\sigma:

x(k+1/2)=x(k)−ηAH(Ax(k)−y)\mathbf{x}^{(k+1/2)} = \mathbf{x}^{(k)} - \eta \mathbf{A}^H(\mathbf{A}\mathbf{x}^{(k)} - \mathbf{y}) x(k+1)=Dσ(x(k+1/2))\mathbf{x}^{(k+1)} = D_\sigma(\mathbf{x}^{(k+1/2)})

where y=Ax+n\mathbf{y} = \mathbf{A}\mathbf{x} + \mathbf{n} is the measurement.

DRUNet (Chapter 27) is ideal for PnP because its noise level map input allows controlling the denoising strength at each iteration.

Example: PnP for Image Deblurring

Use PnP with DRUNet to deblur an image.

Solution
Algorithm
for k in range(n_iter):
    # Data fidelity (gradient step)
    residual = A(x) - y
    x = x - eta * A_adjoint(residual)
    # Denoising step
    sigma = sigma_max * (sigma_min/sigma_max) ** (k/n_iter)
    noise_map = torch.full_like(x[:, :1], sigma)
    x = drunet(torch.cat([x, noise_map], dim=1))
Using Pre-Trained ModelsDomain Adaptation and Transfer