Plug-and-Play Image Restoration with Flow Matching: A Continuous Viewpoint
Fan Jia, Yuhao Huang, Shih-Hsin Wang, Cristina Garcia-Cardona, Andrea L. Bertozzi, Bao Wang
TL;DR
This work addresses the gap between empirical success and theory in plug-and-play flow-based image restoration (PnP-Flow) by deriving a continuous-time stochastic differential equation (SDE) surrogate for PnP-Flow. The SDE view enables principled error quantification, linking restoration accuracy to the Lipschitz constant of the learned vector field and step scheduling, and motivates a Lipschitz-regularized training regime plus an extrapolation-based acceleration. The authors propose IPnP-Flow, an improved method that leverages these insights to achieve superior restoration performance on denoising, deblurring, super-resolution, and inpainting while maintaining similar efficiency to the baseline. Empirical results on CelebA and AFHQ-Cat demonstrate consistent gains over PnP-Flow and other state-of-the-art approaches, validating the practical impact of the SDE-informed design. The work provides a theoretical framework and concrete algorithmic enhancements that advance principled, efficient plug-and-play image restoration with flow matching.
Abstract
Flow matching-based generative models have been integrated into the plug-and-play image restoration framework, and the resulting plug-and-play flow matching (PnP-Flow) model has achieved some remarkable empirical success for image restoration. However, the theoretical understanding of PnP-Flow lags its empirical success. In this paper, we derive a continuous limit for PnP-Flow, resulting in a stochastic differential equation (SDE) surrogate model of PnP-Flow. The SDE model provides two particular insights to improve PnP-Flow for image restoration: (1) It enables us to quantify the error for image restoration, informing us to improve step scheduling and regularize the Lipschitz constant of the neural network-parameterized vector field for error reduction. (2) It informs us to accelerate off-the-shelf PnP-Flow models via extrapolation, resulting in a rescaled version of the proposed SDE model. We validate the efficacy of the SDE-informed improved PnP-Flow using several benchmark tasks, including image denoising, deblurring, super-resolution, and inpainting. Numerical results show that our method significantly outperforms the baseline PnP-Flow and other state-of-the-art approaches, achieving superior performance across evaluation metrics.