Plug-and-Play Image Restoration with Flow Matching: A Continuous Viewpoint

TL;DR

This paper develops a continuous stochastic differential equation model for plug-and-play flow matching in image restoration, providing theoretical insights that enhance performance and accelerate existing methods across various restoration tasks.

Contribution

It introduces a continuous SDE framework for PnP-Flow, enabling error quantification and acceleration techniques, which improve image restoration results.

Findings

Significant performance improvements over baseline PnP-Flow.

Effective error quantification for better step scheduling.

Accelerated restoration via SDE-based extrapolation.

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.