Multi Task Denoiser Training for Solving Linear Inverse Problems

TL;DR

This paper introduces a multi-task training approach for denoisers used in solving linear inverse problems, resulting in improved performance and efficiency across various tasks by fine-tuning within the iterative solver framework.

Contribution

It proposes end-to-end multi-task training of denoisers integrated into inverse problem solvers, enhancing versatility and performance.

Findings

Average PSNR improved by +1.34 dB across six inverse problems

Reduced number of iterations needed for convergence

Denoiser's objective shifts towards better prior gradient approximation

Abstract

Plug-and-Play Priors (PnP) and Regularisation by Denoising (RED) have established that image denoisers can effectively replace traditional regularisers in linear inverse problem solvers for tasks like super-resolution, demosaicing, and inpainting. It is now well established in the literature that a denoiser's residual links to the gradient of the image log prior (Miyasawa and Tweedie), enabling iterative, gradient ascent-based image generation (e.g., diffusion models), as well as new methods for solving inverse problems. Building on this, we propose enhancing Kadkhodaie and Simoncelli's gradient-based inverse solvers by fine-tuning the denoiser within the iterative solving process itself. Training the denoiser end-to-end across the solver framework and simultaneously across multiple tasks yields a single, versatile denoiser optimised for inverse problems. We demonstrate that even a simple baseline model fine-tuned this way achieves an average PSNR improvement of +1.34 dB across six diverse inverse problems while reducing the required iterations. Furthermore, we analyse the fine-tuned denoiser's properties, finding that its optimisation objective implicitly shifts from minimising standard denoising error (MMSE) towards approximating an ideal prior gradient specifically tailored for guiding inverse recovery.