Image Restoration via Primal Dual Hybrid Gradient and Flow Generative Model
Ji Li, Chao Wang
TL;DR
This work extends Plug-and-Play image restoration by embedding flow-matching priors into a Primal-Dual Hybrid Gradient (PDHG) framework. It generalizes data fidelity from the conventional squared $\ell_2$ loss to $\ell_1$ and $\ell_2$ losses, enabling robust handling of Poisson and impulse noise, while using a flow-based denoiser as the implicit proximal operator for the regularizer. The method combines a time-dependent denoiser with adaptive stepsizes and Moreau-based proximal updates to support non-smooth fidelities, and extends naturally to score-based diffusion models. Empirical results on denoising, deblurring, super-resolution, and inpainting show improved performance over Gaussian-noise-oriented PnP methods, with notable gains when using $\ell_1$ fidelity for impulse noise and $\ell_2$ fidelity for Poisson noise. The work offers a practical, simulation-free pathway to leverage powerful flow priors for a broad class of imaging inverse problems with non-Gaussian noise in real-world settings.
Abstract
Regularized optimization has been a classical approach to solving imaging inverse problems, where the regularization term enforces desirable properties of the unknown image. Recently, the integration of flow matching generative models into image restoration has garnered significant attention, owing to their powerful prior modeling capabilities. In this work, we incorporate such generative priors into a Plug-and-Play (PnP) framework based on proximal splitting, where the proximal operator associated with the regularizer is replaced by a time-dependent denoiser derived from the generative model. While existing PnP methods have achieved notable success in inverse problems with smooth squared $\ell_2$ data fidelity--typically associated with Gaussian noise--their applicability to more general data fidelity terms remains underexplored. To address this, we propose a general and efficient PnP algorithm inspired by the primal-dual hybrid gradient (PDHG) method. Our approach is computationally efficient, memory-friendly, and accommodates a wide range of fidelity terms. In particular, it supports both $\ell_1$ and $\ell_2$ norm-based losses, enabling robustness to non-Gaussian noise types such as Poisson and impulse noise. We validate our method on several image restoration tasks, including denoising, super-resolution, deblurring, and inpainting, and demonstrate that $\ell_1$ and $\ell_2$ fidelity terms outperform the conventional squared $\ell_2$ loss in the presence of non-Gaussian noise.