PnP-Flow: Plug-and-Play Image Restoration with Flow Matching (2410.02423v2)

Abstract: In this paper, we introduce Plug-and-Play (PnP) Flow Matching, an algorithm for solving imaging inverse problems. PnP methods leverage the strength of pre-trained denoisers, often deep neural networks, by integrating them in optimization schemes. While they achieve state-of-the-art performance on various inverse problems in imaging, PnP approaches face inherent limitations on more generative tasks like inpainting. On the other hand, generative models such as Flow Matching pushed the boundary in image sampling yet lack a clear method for efficient use in image restoration. We propose to combine the PnP framework with Flow Matching (FM) by defining a time-dependent denoiser using a pre-trained FM model. Our algorithm alternates between gradient descent steps on the data-fidelity term, reprojections onto the learned FM path, and denoising. Notably, our method is computationally efficient and memory-friendly, as it avoids backpropagation through ODEs and trace computations. We evaluate its performance on denoising, super-resolution, deblurring, and inpainting tasks, demonstrating superior results compared to existing PnP algorithms and Flow Matching based state-of-the-art methods.

• The paper introduces a novel image restoration approach by integrating a pre-trained flow matching model within a PnP forward-backward splitting framework.
• It details a method using a time-dependent denoiser and reprojection onto flow trajectories to balance data fidelity with denoising performance.
• Extensive experiments on datasets like CelebA and AFHQ-Cat show that PnP-Flow outperforms existing methods in PSNR and SSIM metrics.

PnP-Flow: Plug-and-Play Image Restoration with Flow Matching

This paper introduces a novel Plug-and-Play (PnP) algorithm called PnP Flow Matching for image restoration tasks. The method leverages the strengths of pre-trained Flow Matching (FM) models within an optimization framework, offering a computationally efficient and memory-friendly alternative to existing approaches.

Methodology

The core idea is to define a time-dependent denoiser using a pre-trained FM model, which is then integrated into a Forward-Backward Splitting (FBS) PnP framework. The algorithm alternates between three key steps: a gradient descent step on the data-fidelity term, a reprojection step onto the learned FM path, and a denoising step using the time-dependent denoiser. The time-dependent denoiser is defined as $D_t = Id + (1-t)v_t^\theta$, where $v_t^\theta$ is the learned velocity field from the FM model. This design is motivated by the fact that $D_t$ approximates the conditional expectation $\mathbb{E}[X_1 | X_t = x]$, where $X_1$ is a sample from the target distribution and $X_t$ is a point along the flow path.

A key aspect of the method is the reprojection step, where iterates are reprojected onto flow trajectories via linear interpolation. Given a current iterate $x$ and a time $t$, the reprojected point is computed as $\tilde{z} = (1-t)\epsilon + tz$, where $\epsilon$ is a noise sample drawn from the latent distribution $P_0$ and $z$ results from a gradient step on the data-fidelity term. This step ensures that the input to the denoiser lies within the support of the flow path, improving the effectiveness of the denoising operation.

Figure 1: Our method on a 2D denoising task ($\sigma=1.5$) with Gaussian distributions.

Implementation Details

The algorithm's performance is sensitive to the choice of the learning rate. The authors suggest using a time-dependent learning rate of the form $\gamma_t = (1-t)^\alpha$, where $\alpha \in (0,1]$. This helps balance the contributions of the data-fidelity term and the denoiser, preventing the algorithm from simply returning the noisy input at later time steps.

The paper also presents a convergence result, stating that if the sequence of iterates produced by the algorithm is bounded and the time sequence $(t_n)_{n \in N}$ satisfies $\sum_{n=0}^\infty (1-t_n) < +\infty$, then the sequence converges.

Experimental Results

The authors conducted extensive experiments on image denoising, deblurring, inpainting, and super-resolution tasks, using the CelebA and AFHQ-Cat datasets. The results demonstrate that PnP-Flow Matching consistently outperforms state-of-the-art FM-based and PnP methods in terms of PSNR and SSIM metrics. Furthermore, the method exhibits stability across different tasks, unlike some competing approaches that perform well on certain tasks but struggle on others.

For instance, on the CelebA dataset, PnP-Flow achieved a PSNR of 32.45 dB on the denoising task, outperforming methods like OT-ODE (30.50 dB) and Flow-Priors (29.26 dB). Similarly, on the super-resolution task, PnP-Flow achieved a PSNR of 31.49 dB, surpassing OT-ODE (31.05 dB) and other baselines. Similar trends are seen with SSIM scores.

Figure 2: Results for random inpainting using PnP-Flow across different iterations (time steps) with corresponding PSNR values.

Advantages

The method offers several advantages over existing approaches:

• It is computationally efficient and memory-friendly, as it avoids backpropagation through ODEs and trace computations.
• It is simple to implement and requires few hyper-parameters.
• It supports different latent distributions and flexible initialization.
• It delivers strong performance across various inverse problems.

Limitations

The authors note that the reconstructions produced by PnP-Flow Matching tend to be slightly over-smoothed, which they attribute to the denoising operation acting as a minimum mean squared error estimator.

Implications and Future Work

PnP-Flow Matching offers a promising framework for integrating generative models into image restoration tasks. The method's computational efficiency and versatility make it a valuable tool for practitioners working on real-world imaging problems.

Future research directions include exploring the use of PnP-Flow Matching for other types of measurement noise, such as Poisson noise, and investigating the use of different latent distributions to model categorical data. Also, further investigation into better theoretical convergence bounds would be useful.

Conclusion

The paper presents a compelling approach to image restoration by combining PnP methods with Flow Matching. The proposed PnP-Flow Matching algorithm demonstrates strong empirical performance and offers several practical advantages over existing techniques. The method's versatility and efficiency make it a valuable contribution to the field of computational imaging.