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

Authors: Ségolène Tiffany Martin, Anne Gagneux, Paul Hagemann, Gabriele Steidl
Primary area: generative models
Keywords: Plug-and-Play, Flow Matching, image restoration, inverse problems, generative modeling

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. Code available at https://github.com/annegnx/PnP-Flow.

Decision: accept (poster)
Reviewer scores: 3, 5, 6, 8

Reviews and discussion:
<thread>
creation date:
1730647196328

thread_type:
official_review

<thread_object>
creator: reviewer_vjzf

<content>
summary:
In this paper, the authors proposed a plug-and-play image restoration method based on flow matching. The reformulation starts from the forward-backward splitting algorithm, where the proximal step is replaced by a denoising step to form the plug-and-play forward-backward splitting algorithm. The authors insert a specific flow matching method, namely straight-line flows into the PnP-FBS framework due to the computation efficiency of the straight-line flows. Formally, the PnP flow matching algorithm consists of three steps: a gradient step on the data fidelity term, an interpolation step, and a PnP denoising step that is specifically designed to denoise inputs drawn from the straight path.

strengths:
1. A new plug-and-play method based on flow matching is proposed in this paper.

2. The paper is well-written.

3. The derivations in this paper are rigorous.

4. The computational complexity and memory footprint of the proposed method is lower than the previous methods due to the careful design.

weaknesses:
1. The restored images seem to be over-smoothed.

questions:
1. Please explain why the computational complexity and memory footprint of the proposed method is lower than the previous method. Is it due to the design of the model or the choice of the straight-line flow?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730692103892

thread_type:
official_review

<thread_object>
creator: reviewer_fvxi

<content>
summary:
This paper proposes to use flow matching in the plug-and-play framework for image restoration. The key is to use FM model as the denoisier. To avoid the numerical challenges, it integrates the implicit FM prior into a custom denoisier.

strengths:
1, It proposes a design a time-dependent denoiser based on a pre-trained velocity field v learned through Flow Matching

2, This denoiser is integrated into an adapted Forward-Backward Splitting PnP framework that cycles through a gradient step on the data-fidelity term, an interpolation step and a denoising step

3, Being computationally efficient and memory-friendly via the use of ODE

weaknesses:
1, Why the percpetual metrics are missing? From the visual results, it also seems that the results tend to be blurry. What’s the underlying reason? Is it due to the gradient step or the interpolation step, or something else? 

2, In addition, one of the advantages of these generative method is its high perceptual quality, but this method seems to have achieved good distortion performance. How about the results of employing the same end-to-end U-Net model as a simple baseline (for example, using the L1 loss)?

3, Can you visualize all the intermidate resutls of all three steps for all time steps? It could better help readers understand the method.

questions:
See weakness
</content>
</thread_object>

</thread>
<thread>
creation date:
1730713683294

thread_type:
official_review

<thread_object>
creator: reviewer_n2r9

<content>
summary:
The paper proposes a zero-shot method (PnP-flow) for Inverse problems based on a pre-trained flow-matching (FM) model. The method combines the plug-and-play (PnP) framework with flow matching  by alternating between gradient descent steps on the data-fidelity term, reprojections onto the learned FM path, and denoising. PnP-flow achieves state-of-the-art (SOTA) results compared to existing PnP and flow-based algorithms  across different image inverse problems.

strengths:
1. The method is training-free which makes it computationally practical.
2. The method achieves SOTA results compared to existing flow-based methods.

weaknesses:
1. My major concern is the lack of comparison to recent zero-shot methods based on a pre-trained diffusion model such as DDNM [1] and DPS [2].
2. The proposed method is non-blind (assume the full knowledge of the degradation model) which limits its applicability.


[1] Wang et al. Zero-Shot Image Restoration Using Denoising Diffusion Null-Space Model. ICLR 2023

[2] Chung et al. Diffusion Posterior Sampling for General Noisy Inverse Problems. ICLR 2023

questions:
1. Could you add comparisons with [1] and [2], or explain why those comparison are missing?
2. Could you comment on the potential applicability/extension of your method to the blind case?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730943735802

thread_type:
official_review

<thread_object>
creator: reviewer_pwbl

<content>
summary:
This paper introduces the PnP-Flow Matching algorithm for addressing imaging inverse problems, including denoising, super-resolution, deblurring, and inpainting. The method combines the Plug-and-Play (PnP) framework with Flow Matching (FM) models by using a time-dependent denoiser to tackle image restoration tasks. Specifically, the algorithm alternates between gradient descent on a data fidelity term, reprojection onto a flow matching path, and denoising. The experiments demonstrate the effectiveness of the proposed method.

strengths:
1. Low memory usage, making it suitable for high-resolution images.
2. Consistently performs well across multiple tasks, showing stable PSNR and SSIM improvements.

weaknesses:
1. The writing quality needs Improvement. Certain explanations lack clarity, particularly in describing the algorithmic process, e.g., the function F.
2. The details of the proposed method are insufficient.
3. The experiment section should be improved. Please refer to the details below.

questions:
1. The formula "y = noisy(Hx)" uses a general definition for the noisy function. It would be helpful if the paper and experiments explored multiple types of noise to assess the method’s robustness.

2. In Tables 1 and 2, the comparison is limited, particularly with only one diffusion-based method, PnP-Diff, which is a workshop paper, not a main conference paper. The authors should include comparisons with more diffusion-based methods, such as DPS, DeqIR, and DDRM, to provide a fuller view of how their method performs relative to the latest diffusion techniques.

3. The authors could enhance the evaluation by including the ImageNet dataset. For the denoising and deblurring tasks. Testing the method across various noise levels and degrees of blur on a large, diverse dataset like ImageNet would offer more insight into how well the algorithm handles different types of degradation.

4. In Figure 3, the visual results do not show a significant improvement over other methods (e.g., in the last row), even though the PSNR scores are higher. 

5. For real-world data with unknown degradation, it would be important to understand how well this method generalizes. 

6. It would strengthen the paper if the authors included examples of failure cases. 

7. In Table 3, not all methods are compared for computational time and memory usage. Including all relevant methods in this comparison would give a clearer picture of how the proposed algorithm stacks up in terms of efficiency across different benchmarks.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732109848935

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General answer to all the reviewers

comment:
We thank all the reviewers for their effort and constructive feedback. All reviewers enjoyed the excellent performance of our method across multiple tasks and its computational efficiency and low memory footprint. 

We uploaded a revised version of the paper according to all concerns of the reviewers with modifications appearing in blue. In particular, we made a big effort to add numerical experiments. We address some common/major points in this general answer.

- Based on the suggestion of reviewer pWBL, we added an experiment incorporating Laplace degradation noise. Benchmarking against other methods demonstrates the adaptability of our PnP-flow method. Here, we outperform competing methods across all metrics.
- We added more diffusion baselines, namely DPS [1] and DDRM [2], as recommended by reviewer n2r9 and pWBL.
- As suggested by reviewers pWBL and n2r9, we added numerical experiments for blind inverse problems where the forward is unknown.
- Following the suggestion of reviewer fVXi, we added a perceptual metric, LPIPS [3].

[1] Chung et al, Diffusion Posterior Sampling for general noisy inverse problems, ICLR 2023

[2] Kawar et al, Denoising diffusion restoration models, NeurIPS 2022

[3] Zhang, The unreasonable effectiveness of deep features as a perceptual metric, CVPR 2018
</content>
</thread_object>

</thread>
<thread>
creation date:
1732530997981

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Friendly reminder

comment:
Dear reviewers, since the discussion period is ending soon, we kindly remind you of giving us feedback on the revised version of our paper. We took great care and efforts in revising it and addressed all your concerns.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733286832540

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Discussion Summary

comment:
The review period is now over, and there unfortunately has not been any discussion regarding our paper, although we really tried to engage one.  We believe that we have solid contributions, and adressed the reviewers concerns. 
Many thanks to the reviewer, who gave already score 8 and underlined again that she/he believes in our work..
</content>
</thread_object>

</thread>


Metareview:
metareview:
Reviewers agree with the practicality of the proposed method and its novelty.
Reviewers pWBL and n2r9 had some concerns about the validation of the model, which the authors addressed and included in their rebuttal. But the reviewer did not further comment on those results.
Reviewer pWBL made further comments about the text clarity and this was again addressed by the author.
Given that the authors responded to reviewer comments and other reviewers were happy with the paper's contribution and evaluation, this paper is recommended for acceptance.