PnP-Flow: Plug-and-Play Image Restoration with Flow Matching
Authors: Ségolène Tiffany Martin
Anne Gagneux
Paul Hagemann
Gabriele Steidl
Venue: ICLR 2025 Poster (2025)
Keywords: Plug-and-Play
Flow Matching
image restoration
inverse problems
generative modeling
Primary area: generative models
TL;DR: We introduce an algorithm combining Plug-and-Play with Flow Matching for solving imaging inverse problems.

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.

Published as a conference paper at ICLR 2025
PNP-FLOW: PLUG-AND-PLAY IMAGE RESTORATION WITH FLOW MATCHING
Ségolène Martin
Anne Gagneux Paul Hagemann
Gabriele Steidl
Technische Universität Berlin
2 ENS de Lyon, CNRS, Université Claude Bernard Lyon 1, Inria, LIP, UMR 5668
* Equal contributions
ABSTRACT
In this paper, we introduce Plug-and-Play (PnP) Flow Matching, algorithm for solving imaging inverse problems. PnP methods leverage the strength of pre- trained denoisers, often deep neural networks, by integrating them in optimiza- tion 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 image sampling yet lack a clear method for efficient use image restoration. We propose to combine the PnP framework with Flow Match- ing (FM) by defining a time-dependent denoiser using a pre-trained FM model. Our algorithm alternates between gradient descent steps the data-fidelity term, reprojections onto the learned FM path, and denoising. Notably, our method is computationally efficient and memory-friendly, 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 utthtpss//giitee
it avoids backpropagation
com annegnx//ggg/////// Pte///nPPPFFooo
INTRODUCTION
Image restoration aims to recover an unknown image Rd from a degraded observation y E Rm
Hr +
where H Fd - 1T is (linear) degradation operator and & describes the underlying additive noise model. Since the problem is usually ill-posed and high dimensional, its treatment challenging. We assume that the image is sampled from random variable X E Rd with density px, and the observation y from a random variable E RT with a density py. Then the maximum posteriori estimator searches for the value with the highest probability of the posterior
arg max =ERd
{ (log (z)) = arg max log Pog (v) + log px (x)},
(1)
ZERd
where the first term of the right hand side the fidelity to the data and the second the prior distribution of the image. Since px is generally unknown, and in the absence of training data, is standard to instead consider regularized optimization problem of the form:
represents
arg min [F(x) + R(x))
(2)
mERd
where F(x) == log log and R: Rd R usually enforces some assumptions on the solu- tion and ensures the existence of a (unique) minimizer. For instance, for Gaussian noise \ o² daLa), the data-fidelity corresponds to F(x) = z || - yl1.. The regularized minimization problem in
(2 be efficiently handled using proximal splitting methods (Boyd et al. 2011 Combettes &
Pesquet 2011).
Plug-and-Pil (PnP) (Venkatakrishnan et al. 2013) methods build upon the insight that the proximal step on the regularization term is effectively denoising operation. They can be performed with non- learned denoisers such BM3D (Chan al. 2016), however recently these were outperformed by

---

Published as a conference paper at ICLR 2025
neural network-based approaches (Meinhardt et al. 2017 Zhang al., 2017 2021 While PnP
methods have demonstrated considerable effectiveness, recent advances in generative models offer more sophisticated framework for learning priors directly from data, surpassing the limitations of hand-crafted or neural denoisers.
In particular, generative models are used as regularizers in Bora al. (2017) Asim al. (2000)
Altekrüger et al. (2023); Wei et al. (2022). Most of these references learn an invertible transport map T between Gaussian latent distribution and the data distribution, using the change of variables formula to obtain the estimated prior px as an explicit function of the known latent density and the map T. In particular the normalizing flow network with the Real-NVP (Dinh et al., 2017) architecture is by default invertible and has an casy-to-evaluate log determinant of the Jacobian of T. However, this architecture yields rather poor results on large images. Therefore, recent work (Ben-Hamu et al. 2024; ZZhang et al.) 2024b; Pokle et al., 2024 Chung et al. 2023 Song et al.)
2023) focuses more on learning diffusion models (Song et al. 202]
Ho et al. [2020] SSoh--Dickstein
Liu et al. 2023), which do not have these architectural constraints and scale better to large images. The main idea is to establish a path between the latent and the target distribution, which can be optimized in simulation-free way.
al. 2015) Flow Matching models (Lipman al. 2023
In this paper, we focus on Flow Matching models, which learn a velocity field v [0, 1] xR" - Rd going from the latent to the data distribution. Once the velocity field learned, sampling from the target distribution can be done by solving an ODE (Chen et al. 2018)
0,, x) = or (tt, f(0, = x,
where æ is sampled from the latent distribution.
However, using Flow Matching to regularize inverse problems such as image restoration tasks is non- trivial. Indeed, when directly solving the MAP problem (T) using the change of variables formula, one faces numerical challenges due to the backpropagation through the ODE. We circumvent this problem by integrating the implicit Flow Matching prior into a custom denoiser that we plug into Forward-Backwanl algorithm.
Contributions. In this paper, we combine PnP methods with Flow Matching and propose a PnP Flow Matching algorithm. Our contributions
follows:
Inspired by optimal transport Flow Matching, we design time-dependent denoiser based on a pre-trained velocity field v learned through Flow Matching.
Forward-Backward This denoiser is integrated into an adapted Forward- Backward Splitting PnP framework, which cycles through gradient step on the data-fidelity term, an interpolation step to reproject iterates onto flow trajectories, and denoising step.
Our algorithm simple implement, requires no backpropagation through the network, and is highly efficient in both memory usage and execution time compared to existing Flow Matching- based methods for solving inverse problems.
We show that, on two datasets and across denoising, deblurring, inpainting, and super- resolution tasks, our method consistently outperforms state-of-the-art Flow Matching- based and PnP methods in terms of both PSNR and SSIM metries. The code with all benchmark methods is available at: and now included the DeepInv library
com/ annegnx/PnP-PF
2 BACKGROUND
We next provide background on both Plug-and-Play algorithms and Flow Matching models.
2.1 PLUG AND PLAY
PnP algorithms were introduced as extensions of proximal splitting methods like Forward- -arkkrrr Splitting (FBS) or the alternating direction method of multipliers (ADMM). These are first-order
http://gtthhb.cmmm/eep////eppppt
2

---

Published as a conference paper at ICLR 2025
latent data
ground truth
Gradient Step
Interpolation Step
PnP Step
1=0
0.56
t= 0.89
Figure 1: Our method on 2D denoising task (a = 1.5) with Gaussian distributions. An OT Flow Matching model is trained to sample from P1 (m, ( ITd), with m 7 and & 0.5. At each time step, it performs a standard gradient step on the datafit, followed by a projection onto flow trajectories at time t, and finally applies the time-dependent denoiser D..
optimization algorithms for solving problem (2). alternating between a proximal step on the regu- larization and gradient and/or proximal step the data-fidelity function. The proximal operator prox,n R2 + R2 > is defined by pro = arg mingeR4 { lly - x 12 RR(x)) In this paper, we assume that the data-fidelity F' is differentiable with Lipschitz continuous gradient and we focus on the FBS algorithm, corresponding to Algorithm| [ The convergence FBS to a minimizer of F R is guaranteed long the stepsize is chosen y E (0, (0,1/Lip(VF)) where Lip(VF) the Lipschitz constant of V (Combettes & Wajs 2005).
Algorithm 1: FBS
Algorithm 2: PnP-FBS
Initialization: (0))) Rd 0 for r 0, 1,
Initialization: (:(0) E R 7 for r = 0, 1,...o do
(r+1) (()) (VF((())
(r+1) x(1) (V(()()
(r+1) prox,, nr (roxx)(((((+1)) (2(++1))
x: [(++1) = D(2(+t1))
PnP step
Observing that computing prox R corresponds solving a Gaussian denoising problem with reg- ularization 7 R2 Venkatakrishnan et al. (2013) proposed to replace it with an off-the-shelf denoiser D Rd Rdd which can be independently designed learned (see Algorithm 2).
2021 12 Sun et al. 2019) Terris et al. (2020) Hertrich al. (2021) 1); Hurault al. (2022a): Tan (2024) neural network denoisers. Most of these PnP-FBS algorithms do not converge (Zhang et al. 2017 Som- merhoff al.) (2019) and convergence guarantees usually assume non con-expansiveness of the denoiser (Pesquet et al. 2021 Ryu et al. 2019 Hurault et al. 2022b). Yet, constraining the Lipschitz con- stant of neural network, e.g., through weight clipping, spectral normalization or averaged operator constructions harms in general its expressiveness (Ryu et al. 2019).
Many PnP-FBS methods, such as Meinhardt et al. (2017); Zhang et al. 2017
2.2 FLOW MATCHING
Let (2(dd) denote the subspace of probability measures Rd having finite second moment. Let Po denote the latent probability measure and Pi the target probability measure (i.e., the data distribution). We denote the push-forward operation. We denote I(Po, P1) the set of couplings P(R Idd) having marginals P, Py (RR), i 0. 1. and I To Pt) the set of optimal cou- plings minimizing the Wasserstein-2 distance (Villani 2008). We now fix coupling E T(Po, Pt).
The core idea behind Flow Matching
define specifie target probability path t Pt, [0, 1), between Po and Pt. Let P. == (4t) where the map e (Zo, 21) = (1 t) + tæ interpolates between latent sample) æn and (a data sample) 21. The path P: can be shown be an absolutely
continuous curve, so there exists Borel vector field [0 x Rd - Rd such that the curve satisfies the continuity equation
O.P. + (Per)
(CE)
3

---

Published as a conference paper at ICLR 2025
in the sense of distributions (Ambrosio et al. 2008 In addition, there exists a solution f [0, 1] x R + IR of the ODE
0ff(t,z)= = (x(((t,)))) f(0,z)
(3)
such that P. f(t, Po. If / is known, then sample x1 from the target distribution P1 can be drawn by first sampling To ~ Po and then defining 21 = f1)
The goal of Flow Matching is to learn the velocity field of the flow ODE (3) This learning process consists in minimizing the loss function
(FM == [0.1 [x)](()) (=()))?
where v is parametrized by neural network with weights 0. Unfortunately, in practice the true velocity field (z( is not available. However, Lipman et al. (2023) showed that minimizing (0) equivalent minimizing the Conditional Flow Matching (CFM)los) loss: LCFM(0) connst, where
(f) == (0- [lle] (ec(2o, 21)) (z1 =0)_l3]
Minimizing this loss only requires sampling from the coupling æ. For example, if IT Po x P1 is the independent coupling, we only need samples from the latent and target distributions. If instead, ı E T. Pt) an optimal coupling, then we can use a standard optimal transport solver to (approximately) sample from as proposed in Pooladian et al. (203)) Tong et al. (2024).
Straight-line flows Let [0, IRd - R and let be solution to the flow ODE associated with v. Given (Xo, X1) ~ т, we call (f, v) straight-line Flow Matching pair connecting X, and
X, if X, (t. Xo) almost surely, where X, is defined as X, e, ex((Xo X1) tX + +1 X)
et 2023 Pooladian et al. 2023). Note that this directly implies or (f Xo)) = (t,
0,X,
Xa almost surely.
We will focus on such straight-line flows later in the paper. Straight or nearly straight paths are preferred since they represent the shortest distance between two points and can be simulated with fewer steps of an ODE solver. This has been explored in many recent works that want to speed up sampling, such as Kornilov et al.] (2024); Lee et al. (2023); Yang et al. (2024). A particular case of straight-line pair is given by OT couplings (Pooladian et al. 2023 Tong et al. (024).. Indeed, for и E [Po, P1) with Monge map T Rd + Rd such that X, T(x)), there exists a velocity
field for which the pair (Pe, v.) is solution of the continuity equation CC and such that for all e Rdd
==(((tx))=T((t)
(5)
where f is the solution of the ODE (3). see Ambrosio (2008) Pooladian al. (2023): Chemsed-
dine et al. (2024) for the precise conditions. By integrating (5). it is clear that (f. is -traight-line
Flow Matching for the coupling (Xo. T (Xo)). since X, (1 XtX, = (1- x
Xo + So (f Xo))ds = Xo).
PNP MEETS FLOW MATCHING
3.1 DENOISING OPERATOR FROM FLOW MATCHING
Let Xo Po and X₁ P. with joint distribution (Xo, w. Assume we have access pre-trained velocity field u.. Then we define, for t E (0, 1). the following time-dependent denoiser
D. = Id + (1 t)of
(6)
and we propose to use it within a PnP framework.
To motivate the choice of the denoiser in (6) recall that for a fixed time [0, 1). the minimizer of of the CFM loss 4 over all possible vector fields reads
1 (x) E(X, XolX. 20),
4

---

Published as a conference paper at ICLR 2025
where X, := e (Xu, X1) (1-t1 Xt + (Benton et al. 2024
Liu et al. 2023) in the ideal case where uy = 2/- Then it follows that, for any I E R" and f E [0,1]
Assume
Dz(x) + (1 t) (2)
(x [x + () Xo)
Ex [X, =x]
Hence, the operator Dr can be understood at the best approximation of X given the knowledge of X5, which is also used in Pokle et al. (2024) and Zhang al. (2024b).
Just as standard PnP denoisers minimize the MSE loss between noisy and clean samples, the operator (Ix)) projecting any noisy point
Dr is the minimizer of L Li-problem min, (1)=xx1===1]11 taken along the path onto the target distribution. In particular, the following proposition holds, demonstrating that the best denoising performance achieved with straight-line flows.
Proposition 1. Assume v := v is continuous and assume that, given v, the Flow ODE ( has unique solution f. Then the denoising loss (x.....))<al [x (X,) X1 is equal to 0 for all t [0, 1]. if and only if the couple (f, v) is straight-line Flow Matching pair between Xo and X1.
The proof can be found in Appendix A.I
As stated in Section 2.2, a special case of straight-line flows is given by OT Flow Matching for which we recover (X,, X1. Indeed, in this case, D, reduces to (Xx)) X,, with T' the Monge map between Po and Pt. We next discuss whether diffusion models induce straight-line flows.
Remark 2 (Flow Matching versus diffuusin models). Contrary to OT Flow Matching, diffusson models, which are also flow ODE methods, do not generally induce straight paths. Indeed, during diffusion training, the target probability path takes the form X, = X, + 3,X, with a E (0, 2021 Liu et al. 2023)), which clearly does not match the desired
and B, / a (Song et al. straight path X, - x)x
L. The non-straightness of the flow generated by diffusion models is illustrated in Liu et al. 2023 Liu (2022, Figure 5). On the other hand, a non-straight path obtained with Flow Matching (for instance, using an independent coupling л) can be rectified
to straighter one following the procedure described in Liu et al. (2023).
3.2 PNP FLOW MATCHING ALGORITHM
In the previous section, we built denoiser D, (defined 6 We now want to plug in a Forward- Backward Splitting algorithm in order to solve inverse problems. Yet, our algorithm differs from the classical PnP-FBS (Algorithm [ in two key aspects. First, the iterations of the our algorithm de- pend time because of the definition of D,. Second, we introduce an intermediate reprojection step between the gradient step on the data- fida-fidelity term and the denoising step. More precisely, at each time t [0, 1]. given the current iterate I, our algorithm does the following updates:
1. Gradient step: gradient step on the data-fidelity term, mapping x to 2 = x - Fx) for a given learning rate 7 > 0.
2. Interpolation step: In standard Forward-Backward scheme, the denoiser operator is applied right after the gradient step. Yet, as diseussed earlier, our operator D, was specifically designed to cor- rectly denoise inputs drawn from the straight path X, = (1 (1-1) X + t.,,- If the output z from the gradient step time does not lie in the support of Xx, there is a high chance that the denoising will not be effective, hence the need to "reproject" along the flow paths be- fore applying D. To achieve this, we perform linear interpolation on z, as illustrated in Figure 2 at time f, we define
Interpoaatoon
(1 - (Je +
PuP denricing step
(1-t)e + tz, where is noise sample drawn from Po. A similar idea can be found in (Cheng et al. 2025). Note that while & is sampled from Po, it is not necessarily coupled to 2 P1 via и. If it were, D, would map ă directly back to 2, annihilating the denoising effect.
Xo
Xi
Figure 2: Illustration of the interpolation step.
3. PnP Denoising step: the operator D, is applied to the output à of the interpolation step, regularizing the current image by pushing towards the distribution of X1.
5

---

Published as a conference paper at ICLR 2025
The resulting soscrete-time algorithm is given in Algorithm Figure illustratee the three steps of the algorithm on a denoising problem with a Gaussian prior.
Algorithm 3: PnP Flow Matching
Input: Pre-trained network v by Flow Matching, time sequence (l either finite with n/N, N E N infinite with lim + too in adaptive stepsizes (7x)n.
Initialize: 20 E Rd
for n 0,1,.... do
2n = In 7n V.n)).
>> Gradient step on the data-fidelity term Linear interpolation b PnP step with denoiser (6)
Zn = tn E Zn
Po
En+1 Dt., (2n)
return 1n++
Remark (Averaging in the denoising step). Instead of drawing one noise realization, we can also average over multiple samples ~ Po in the last step of the algorithm: In+1 == EE [Dr, (Zn (+))] with (e) == b, + 2n The algorithm's 's output is then deterministic. In practice, averaging over few realizations slightly improves the numerical results.
Time dependent learning rate Using constant learning rate independent of time can give too much importance to the data fit. For example, if 7 for all E [0. 1] on simple denoising task, the algorithm will return the noisy sample since Dj = ld. To prevent this, should decrease with t to balance the contributions of the datafit and the denoiser. We set ft (1 t) with (0, 1] for the remainder of the paper. This choice yields good numerical results in our experiments, but alternative values for y could also be explored.
Convergence Assuming that the sequence produced by the algorithm is bounded in the infinite time regime, we have the following convergence result. The proof can be found in Appendix A.2
Proposition 4. Assume that F I + R. is continuously differentiable and that the learned vector field [0, 1] x Red Rd is continuous. Assume Po has bounded support. Let the time sequence (tn)neN satisfy En. (1-t)) < too and let "io == - tm n. E If the sequence (In)net obtained by Algorithm is bounded, then converges.
RELATED WORK
Our approach combining PnP with Flow Matching relates to several existing methods.
Pre-trained Flow Matching methods Using pre-trained Flow Matching models for regularizing image inverse problems has been the focus of several recent works.
OT-ODE (Pokle al., 2024 assumes a Gaussian latent distribution and uses Tweedie's formula to derive a new velocity field (z, y) E[X, - Xo X, x, Y yi from the original velocity field (z(x) E[X; Xo X, æ).. all without requiring retraining. In practice, they sample from the posterior distribution X1 Y by solving the ODE with the new velocity field using an Euler scheme.
In Zhang al. (024b), the authors introduce Flow-Priors, a method tackle the MAP problem by approximating it as sequence of time-dependent MAP subproblems. Using Tweedie's 's formula, they show that for < the gradient of the distribution P. of X, can be computed in closed form, allowing for the use of gradient descent to optimize these subproblems. However, the closed-form expression for the gradient relies on the assumption an independent coupling æ and a Gaussian latent distribution. Besides, their method requires computing Tr V which is expensive.
In Ben-Hamu et al. (2024). an implicit regularization approach called D-Flow considered: instead of minimizing the Gaussian data-fidelity function æ - H yl they minimize the latent loss 2 + (f((1,))) 2)) y where ff is a solution to the flow ODE given the pre-trained network. The two problems loss not convex, first- or second-order optimization methods may not find the global minimizer. Interestingly, this is beneficial because the true minimizer of the original problem is simply the
theoretically equivalent since f1. is invertible. However, since the latent

---

Published as a conference paper at ICLR 2025
useudo-inverse, which may not be desirable. The authors optimize the latent loss by backpropagating through the ODE solution with a few Euler steps, though this remains computationally expensive.
(PnP) Difftusion methods While we present the first PnP method based on Flow Matching, related works combine diffuusin models with the PnP framework. We refer to Appendix A.12 for further discussion.
In Zhu al. (2023), a Half Quadratic Splitting algorithm is used, alternating between a proximal step on the data-fidelity term and proximal step on the regularization. Following the PnP strategy, the proximal step for the regularization term is replaced with denoising step using pre-trained diffusion model. The denoiser they use is reminiscent of the one we use, where the velocity is replaced by the gradient of the score function. Their method also includes an interpolation step with random noise, mapping the estimated image at each iteration back to the diffusion path. In Garber & Tirer 2024), the proximal step on the fata-fidelity term in Zhu al. (2023) is traded for gradient step. Finally, conditional image restoration is explored in (Graikos et al. (022), which uses a more relaxed definition of "plug and play". They integrate a pre-trained diffusion model under different conditions, leading to a variational objective similar to methods like Mardani et al. (2024).
NUMERICAL EXPERIMENTS
5.1 BASELINES
We benchmark our method against three state-of-the-art Flow Matching-based restoration meth- ods: OT-ODE (Pokle et al. 2024). D-Flow (Ben-Hamu et al., 2024) and Flow Priors (Zhang et al., 20246). As no official implementations were publicly available for these methods, our own based on the descriptions provided in their respective publications. We have made every effort ensure faithful implementations, included in the code attached to this paper. We also bench- mark our method against PnP-Diff (Zhu et al., 2023). a PnP algorithm based on diffusion models. Additionally, we compare our approach with the state-of-the-ar PnP-FBS (Hurault et al.) 2022a).
developed
5.2 EXPERIMENTAL SETUP
2015). with images resized Datasets We evaluate all methods on two datasets: CelebA YYagg et al. to 128 x 128, and AFHQ-Cat, a subset of the Animal FacesHQ dataset (Choi et al. 2020) focused on the cat class, with images resized to 256 x 256. All images are normalized to the range [-1, 1]. For CelebA, we use the standard training, validation, and test splits. For AFHQ-Cat, as no validation split is provided, we randomly select 32 images from the test to create validation set.
Models For each dataset, trained a Flow Matching model from scratch using the Mini Batch OT Flow Matching approach (Tong et al.) 2024), as this choice of coupling usually leads to straight paths. We used standard Gaussian the latent distribution and U-Net (Ronneberger al.) 2015) taken from Huang et al. I): Ho et al. (220) as the model. The training parameters were learning rate of 10 200 epochs with batch size of 128 for CelebA, and 400 epochs with a batch size of 64 for AFHQ-Cat. We train the denoising network for the PnP method PnP-GS (Hurauli al. 2022a) employing the same U-Net architecture. For PnP-Diff, because training a diffusion model with the same U-Net architecture as Flow-Matching yielded poor results due to insufficient number parameters, we used the pre-trained model from Choi et al. (202T), implemented in the DeepIny library (Tachella et al. 2023). Note that the pre-trained diffusion model was trained on the FFHQ dataset (Krrras et al. 2019), making the comparison indirect, but we had no alternative.
Settings for the experiments We evaluate the methods using 100 test images across five restora- tion problems: denoising with Gaussian noise (a 0.2); deblurring using 61 x 61 Gaussian kernel (o& = 1.0 for CelebA, dg = 3.0 for AFHQ-Can super-resolution (2x downsampling for CelebA, 4x for AFHQ-Cat); box- vox-inpainting with centered x mask (s 40 for CelebA, AFHQ-Cat); and random pixel inpainting (70% masked pixels). For deblurring, super-resolution, and box-inpainting, add Gaussian noise with o 0.05, and for random inpainting a 0.01.
80 for
7

---

Published as a conference paper at ICLR 2025
We optimize the hyper- ppeer-paramterrs for each method using a grid search on the validation set, selecting the configuration that yields the highest Peak Signal-to-Nor to-iise Ratio (PSNR). detailed in Appendix A.8
Hyper-parameters
The optimal values identified for each dataset and problem scenario Our proposed method has two hyper-parameters: the exponent a in the learning rate schedule ?nn. (1 - 1.-)) 1n and the number of uniformly spaced time steps was set to N 100 for most experiments.
We averaged the results of the denoising step over realizations of the interpolation step.
5.3 MAIN RESULTS
We report benchmark results for all methods across several restoration tasks, measuring average PSNR and Structural Similarity (SSIM) on 100 test images. To ensure reproducibility, all experi- ments were seeded. Results are presented in Table | for CelebA and Table for AFHO-Cat. "N/A" indicates cases where the method is inapplicable; for instance, PnP-GS (Hurault et al. [202aa is not designed for generative tasks like box inpainting, and PnP-Diff relies on a diffusion model trained another dataset, making its evaluation box inpainting inappropriate.
The tables show that our method consistently ranks first or second in both reconstruction metrics across all tasks and datasets. More importantly, demonstrates stability across tasks, unlike other methods. For example, D-Flow and Flow-Priors perform well in box inpainting but struggle with denoising, while PnP-GS, PnP-Diff, and OT-ODE excel in denoising and deblurring but perform worse on pseudo-generative tasks.
In terms of visual quality (Fig. 3 Fig. 4 and Appendix A.7 our method produces realistic, artifact- free images, though sometimes slightly over-smoothed. While D-Flow generates realistic images, it occasionally suffers from hallucinations (e.g., eye color shifts in CelebA denoising tasks). Flow- Priors introduces noise and artifacts, while OT-ODE captures textures well but struggles with image generation. Finally, Appendix A.6 shows the progression of the reconstruction with respect to time.
Table I: Comparisons of state-of-the-art methods on different inverse problems on the dataset CelebA. Results are averaged across 100 test images.
Denoising 0.2
Dehlurring
Super-res. 0.05, x2
Rand, inpaint. 0.01,
Box inpaint.
Method
0.05,
1.0
70%
0.05, 40 40
PSNR
SSIM
PSNR
SSIM
PSNR
SSIM
PSNR
SSIM
PSNR
SSIM
Degraded
20.00
0.348
27.67
0.740
10.17
0.182
11.82
0.197
22.12
0.742
PnP-Diff
31.00
0.883
32.49
0.911
31.20
0.893
31.43
0.917
N/A
N/A
PnP-GS
32.45
0.908
33.65
0.924
30.69
0.889
28.45
0.848
N/A
N/A
OT-ÔDE
30.50
0.867
32.63
0.915
31.05
0.902
28.36
0.865
28.84
0.914
D-Flow
26.42
0.651
31.07
0.877
30.75
0.866
33.07
0.938
29.70
0.893
Flow-Priors
29.26
0.766
31.40
28.35
0.717
32.33
0.945
29.40
0.858
PnP-Flow (ours)
32.45
0.911
34.51
0.940
31.49
0.907
33.54
0.953
30.59
0.943
Table 2: Comparisons of state-of-the-art methods on different inverse problems on the dataset AFHQ-Cat. Results are averaged across 100 test images.
Denoising
Deblurrigg 0.05, a 3.0
Super-res. 0.05, x4
Rand. inpaint. 0.01,
Box inpaint.
Method
0.2
70%
80
PSNR
SSIM
PSNR
SSIM
PSNR
SSIM
PSNR
SSIM
PSNR
SSIM
Degraded
20.00
0.319
23.77
0,514
11.59
0.216
13.35
0.234
21.50
0.744
PnP-Diff
30.27
0.835
27.97
0.764
23.22
0.601
31.08
0.882
N/A
N/A
PnP-GS
32.34
0.895
27.33
0.749
21.86
0.619
29.61
0.855
N/A
OT-ÔDE
29.90
0.831
26.43
0.709
25.17
0.711
28.84
0.838
23.88
0.874
D-Flow
26.22
0.620
27.49
0.740
24.10
0.595
31.37
0.888
26.69
0.833
Flow-Priors
29.32
0.768
25.78
0.692
23.34
0.573
31.76
0.909
25.85
0.822
PnP-Flow (ours)
31.65
0.876
27.62
0.763
26.75
0.774
32.98
0.930
26.87
0.904
8

---

Published as a conference paper at ICLR 2025
Clean
Degruded
PnP-Diff
PaP-GS
OT-ODE
D-Flow
Flow-Priurs
PnP-Plow
FoNR 2003
FOND
1161
F5NR
PANE 3E.
PSNIE 2017
PUNIC 11.07
1SN00: 50.32
ISSHE 3137
IANNE 32.20
PONI 15
PASAL 1252
PANKL 20.2
PANK 2035
PONK 2431
PANE 20.10
PSNE 35.00
PANK 25.5
NIA
PENIL
PENIL 3003
PENK
PESSL 32.30
PSNX
Comparison of image restoration methods CelebA: denoising (Ist row), Gaussian deblurring (2nd row), super-resolution (3rd row), random pixel inpainting (4th row), box-inpainting (5th row). N/A means "method not applicable" Zoom in to that PnP-Flow performs consistently well across all tasks compared to the baselines.
Figure 3:
Clean
Degraded
PnP.Diff
PnP-GS
OT-ODE
D-Flow
Flow- Priors
PaP-Flow
15NI 1003
RSME 2011
PANE: 10
PSNP 10.70
PANP 76
FSNP 71.1
PaNe 15.33
FANA
PANA:
PING
TRNN 20.1
T2NT 75.53
FANe 24
PANKL 15:37
PNNIL 2U3T
PSNR. 11.42
13N0: 20.30
2331 30.00
P5N0 10.30
N/A
N/A
Figure 4: Comparison of image restoration methods on AFHQ-Cat: denoising (Ist row), Gaussian deblurring (2nd row), super-resolution (3rd row), random pixel inpainting (4th row), box box-inpainting (5th row). N/A means "method not applicable".

---

Published as a conference paper at ICLR 2025
5.4 PRACTICAL ASPECTS
Time and memory metries per image on the deblurring task on CelebA (128 128).
Table
Computation time and memory All experiments in this section are conducted on single NVIDIA RTX 6000 Ada Generation with 48GB RAM. We measure the averaged time per image to complete deblurring task on the CelebA dataset. We also compute the peak GPU memory load per image. The results are averaged over 100 images (25 batches of 4 images each). We use the same settings as those used for reporting perfor- mance metrics on CelebA. Results are in Table 3
Computation time
GPU peuk
Method
memory load
OT-ODE
1.50s
0.65GB
Flow-Priors
16.01s
2,96GB
D-Flow
32.19s
5,91GB
PnP-Flow (ours)
3.40s
0.10GB
Sensitivity to initialization Notably, our algorithm does not rely on a good initialization. By starting the algorithm at time which is then given to the denoiser. As a result, the algorithm's 's performance is independent of the initialization. In Appendix A.4 we stress that this not the case for other methods.
0, the linear interpolation step initially outputs pure noise,
Choice of the latent distribution Our algorithm can be used with any latent distribution. This in the spirit of Flow Matching models (Lipman al., 2023). which do not rely on a Gaussian latent as opposed to diffusson models. In particular, recently there has been trend of modelling categorical data with Flow Matching (Boll et al. 2024 Stark et al. 2024: Davis et al. 2024). Our method does not rely Gaussian latent, contrary to the other Flow Matching restoration methods. As an example to illustrate the performance of our approach in non Gaussian case, in Appendix A.3 we conduct an experiment with Dirichlet latent distribution, inspired by (Stark et al. 2024).
Adaptability to any straight-line flows Our denoiser is rooted in the straight-line flow frame- work, which motivates our use of OT Flow Matching. However, other choices of FM models are possible. Notably, Rectified Flows (Liu et al.) (003) are of particular interest, as the method de- scribed allows for the straightening of any flow model. In Appendix AA.5] we show how our method performs similarly to what we observed in Section 5.3 using pre-trained Rectified Flows.
6 CONCLUSION
We introduced PnP-Flow Matching, and compared it to recent Flow Matching and PnP methods. A great strength of our method is versatility: it requires few hyperparameters, uses minimal memory and computational resources, delivers very good performance across various inverse problems, and supports different latent distributions as well as flexible initialization. Regarding limitations, our re- constructions seem to be more the smooth side, which relates to the fact that our denoiser is mean square estimator. This however is a common tradeoff in regularization of inverse problems (Blau) & Michaeli 2018). Next, would be valuable to apply our method other types of measurement noise, such as Poisson noise (Hurault et al.) (003). Preliminary results with Laplace noise are pro- vided Appendix
Furthermore, our method could be explored in the context of blind inverse problems, as suggested by the promising results in Appendix A4 Another potential improvement would be to leverage different latent distributions (Boll et al. 2024; Gat al.) 2024 Stark et al., 2024 for modeling categorical distributions. In particular, Dirichlet distributions are widely used in biological and molecular data analysis, making this a relevant domain for further investigation. Finally, ambitious question for future research would be to explore how the PnP-flow algorithm could be adapted into posterior sampling algorithm.
10

---

Published as a conference paper at ICLR 2025
Acknowledgements Ségolène Martin's work is carried out in the framework of the DFG funded Cluster of Excellence EXC 2046 MATH+ (project ID: AA5-8). Anne Gagneux thanks the RT IASIS for supporting her work through the PROSSIMO grant. Anne Gagneux also acknowledges support from the Berlin Mathematical School of MATH+, which funded her research visit to TU Berlin. Paul Hagemann acknowledges funding by the German Research Foundation (DFG) within the project SPP 2298 "Theoretical Foundations of Deep Learning". The authors thank the Blaise Pascal Center for the computational means. It uses the SIDUS (Quemener Corvellec 2013) solution.
The authors would like to thank Yasi Zhang for helping with the Flow priors code, Bastian Boll for helpful discussions and Marien Renaud for spotting some mistakes.
Ethics Statement Our method makes use of pre-trained generative models, such that the general concerns of such models apply. In particular, these generative models carry inherent biases and can potentially be misused.
However, our work foundational and we believe there are many relevant applications, such as med- ical imaging or scientific applications which outweigh the negative usages. Furthermore our work involves only small Flow Matching models, which carry very low risk of being used for malicious purposes.
Reproducibility Statement We implemented all the baselines and release the code in the supple- mentary material. In Appendix A we state the hyper- syper-parameter search procedure and the found values. Further, all our theoretical results state the precise assumptions and contain full proofs.
REFERENCES
Fabian Altekrüger, Alexander Denker, Paul Hagemann, Johannes Hertrich, Peter Maass, and Gabriele Steidl. PatchNR: learning from very few images by patch normalizing flow regular- ization. Inverse Problems, 3966)0664006, 2023.
Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient Flows in Metric Spaces and in the Space of Probability Measure. Lectures Mathematics ETH Zürich. Birkhäuser, 2nd edition, 2008
Muhammad Asim, Max Daniels, Oscar Leong, Ali Ahmed, and Paul Hand. Invertible generative models for inverse problems: mitigating representation error and dataset bias. In International Conference on Machine Learning (ICML), 2020.
Heli Ben-Hamu, Omri Puny, Itai Gat, Brian Karrer, Uriel Singer, and Yaron Lipman. D-Flow: Differentiating through flows for controlled generation. In International Conference Machine Learning (ICML), 2024.
Joe Benton, George Deligiannidis, and Arnaud Doucet. Error bounds for flow matching methods. Transactions on Machine Learning Research, 2024.
Yochai Blau and Tomer Michaeli. The perception-distortion distortion tradeoff. IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR). IEEE, 2018
Bastian Boll, Daniel Gonzalez-Alvarade Stefania Petra, and Christoph Schnörr. Generative signment flows for representing and learning joint distributions discrete data. arXiv preprint arXiv:2406.0m 04527, 2024.
Ashish Bora, Ajil Jalal, Eric Price, and Alexandros G Dimakis. Compressed sensing using genera- tive models. In International Conference on Machine Learning (ICML), pp. 537-546, 2017.
Stephen Boyd, Neal Parikh, Eric Chu, Borja Peleato, Jonathan Eckstein, et al. Distributed optimiza- tion and statistical learning via the alternating direction method of multipliers. Foundations and Trendy@ in Machine learning, 3112, 2011.
Stanley H Chan, Xiran Wang, and Omar A Elgendy. Fixed-point convergence and applications. IEEE Transactions on Computational Imaging, 3(1); 84-98, 2016.
admm for image restoration:
11

---

Published as a conference paper at ICLR 2025
Jannis Chemseddine, Paul Hagemann, Gabriele Steidl, and Christian Wald. Conditional Wasserstein distances with applications in Bayesian OT flow matching, arXiv preprint arXiv:2024, 2024.
Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David K. Duvenaud. Neural ordinary differential equations. Advances Neural Information Processing Systems (NeurlPSS, 2018.
Chaoran Cheng, Boran Han, Danielle C. Maddix, Abdul Fatir Ansari, Andrew Stuart, Michael W.
Mahoney, and Bernie Wang. Gradient-free generation for hard-constrained systems. In Interna- tional Conference on Learning Representations (ICLR), 2025. URL
Jooyoung Choi, Sungwon Kim, Yonghyun Jeong, Youngjune Gwon, and Sungroh Yoon. Ilvr: Con- ditioning method for denoising diffusion probabilistic models. In CVF International Conference on Computer Vision (ICCV), volume I, pp. 2, 2021.
Yunjey Choi, Youngjung Uh, Jaejun Yoo, and Jung-Woo for multiple domains. IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 8188-8197 2020.
Woo Ha. StarGAN v2: Diverse image synthesis
Hyungjin Chung, Jeongsol Kim, Michael Thompson Mccann, Mare Louis Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. In International Conference on Learning Representations (ICLR), 2023.
Patrick L. Combettes and Chrrsthistee Pesquet. Proximal splitting methods in signal process- ing. In Fixed-point Algorithms for Inverse Problems in Science and Engineering, pp. 185-212. Springer, 2011.
Patrick L Combettes and Valérie R Wajs. Signal recovery by proximal forward-backward splitting. Multiscale Modeling Simulation, 4(4):1168 11000100000, 2005.
Oscar Davis, Samuel Kessler, Mircea Petrache, ismail Îlkan Ceylan, Michael Bronstein, and Avishek Joey Bose. Fisher flow matching for generative modeling discrete data. arXiv preprint Xrv::22405. 146644 2024.
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real NVP. In International Conference on Learning Representations (ICLR), 2017.
Tomer Garber and Tom Tirer. Image restoration by denoising diffusion models with iteratively
preconditioned guidance. In IEEF/CVF Conference Computer Vision and Pattern Recognition (CVPR), pp. 25245-25254, 2024.
Itai Gat, Tal Remez, Neta Shaul, Felix Kreuk, Ricky T. Q. Chen, Gabriel Synnaeve, Yossi Adi, and Yaron Lipman. Discrete flow matching. arXiv preprint arXiv:2407.15595 2024.
Alexandros Graikos, Nikolay Malkin, Nebojsa Jojic, and Dimitris Samaras. Diffusion models as plug-and-play priors. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems (NeurIPS), 2022.
Jingyu Guo. Hongliang Qi, Yuan Xu, Zijia Chen, Shulong Li, and Linghong Zhou. Iterative image reconstruction for limited-angle CT using optimized initial image. Computational and Mathe- matical Methods in Medicine, 2016(1):583640, 2016.
Yutong He, Naoki Murata, Chieh-Hsin Lai, Yuhta Takida, Toshimitsu Uesaka, Dongjun Kim, Wei- Hsiang Liao, Yuki Mitsufuji. J Zico Kolter, Ruslan Salakhutdinov, and Stefano Ermon. Manifold preserving guided diffusion. In The Twelfth International Conference on Learning Representa- tions, 2024. U httttss:///peeneeee net.iiewwwweeeeeeeeeeeeeeeeeeeee wwtp:/eeeteettttteteeeeet
Johannes Hertrich, Sebastian Neumayer, and Gabriele Steidl. Convolutional proximal neural net- works and aing-and-play algorithms. Linear Algebra and its Applications, 631:203-2234 2021.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Infor- mation Processing Systems (NeurIPS), volume 33, pp. 6840-6851, 2020.
12

---

Published as a conference paper at ICLR 2025
Chin-Wei Huang, Jae Hyun Lim, and Aaron C Courville. A variational perspective diffusion- based generative models and score matching. In M. Ranzato, A. Beygelzimer, Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems (NeurIPS), volume 34, pp. 22863-22876, 2021.
Samuel Hurault, Arthur Leclaire. and Nicolas Papadakis. Gradient step denoiser for convergent plug-and-play In International Conference on Learning Representations (ICLR), 2022a.
Samuel Hurault, Arthur Leclaire, and Nicolas Papadakis. Proximal denoiser for convergent plug- and-play optimization with nonconvex regularization. In International Conference on Machine Learning (ICMI.), 2022b.
Samuel Hurault, Ulughek Kamilov, Arthur Leclaire, and Nicolas Papadakis. Convergent bregman plug-and-play image restoration for poisson inverse problems. In A. Oh, T. Naumann, A. Glober- son, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Information Processing Systems (NeurIPS), volume 36, pp. 27251-27288 2023.
Tero Karras, Timo Aila, Samuli Laine, and Jaakko Lehtinen. Progressive growing of gans for im- proved quality, stability, and variation, In International Conference on Learning Representations (ICLR), 2018.
Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 4401 4410, 2019.
Bahjat Kawar, Michael Elad, Stefano Ermon, and Jiaming Song. Denoising diffusion restoration models. In Advances in Neural Information Processing Systems, 2022.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 0rXiv:1412.6998. 69800 2017.
Nikita Koriloov, Alexander Gasnikov, and Alexander Korotin. Optimal flow matching: Learning straight trajectories just one step. arXiv preprint 17Xiv:2403./311 2024.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):2278 2324, 1998. doi: 10.1 726791.
Sangyun Lee, Beomsu Kim, and Jong Chul Ye. Minimizing trajectory curvature of ODE-based generative models. In International Conference Machine Learning (ICML), 2023.
Anat Levin, Yair Weiss, Fredo Durand, and William T Freeman. Understanding and evaluating blind deconvolution algorithms. In 2009 IEEE conference computer vision and pattern recognition, pp. 1964-1971. IEEE, 2009.
Xiang Li, Soo Min Kwon, Ismail R. Alkhouri, Saiprasad Ravishankar, and Qing Qu. Decoupled data consistency with diffusion purification for image restoration. arXiv preprint arXiv:2403. 06054, 2024.
Yaron Lipman, Ricky Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In International Conference Learning Representations (ICLR), 2023.
Qiang Liu. Rectified flow: A marginal preserving approach to optimal transport. arXiv preprint arXiv:2209.14577 2022.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning generate and transfer data with rectified flow. In International Conference on Learning Representations (ICLR), 2023.
Raunak Manekar, Zhong Zhuang, Kshitij Tayal, Vipin Kumar, and Ju Sun. Deep learning initialized phase retrieval. In Neural Information Processing Systems (NeurIPS) Workshop, 2020.
13

---

Published as a conference paper at ICLR 2025
Morteza Mardani, Jiaming Song, Jan Kautz, and Arash Vahdat. A variational perspective solving inverse problems with diffusion models. In International Conference on Learning Representa- tions (ICLR), 2024.
Tim Meinhardt, Michael Moeller, Caner Hazirbas, and Daniel Cremers. Learning proximal opera- tors: Using denoising networks for regularizing inverse imaging problems. In IEEE International Conference on Computer Vision (ICCV), pp. 1799-1808. IEEE Computer Society, 2017.
Cean-Christophe Pesquet, Audrey Repetti, Matthieu Terris, and Yves Wiaux. Learning maximally monotone operators for image recovery. SIAM Journal on Imaging Sciences, (13)):106-12333
2021.
Ashwini Pokle, Matthew J. Muckley, Ricky T. Q. Chen, and Brian Karrer. Training-free linear image inverses via flows. Transactions Machine Learning Research, 2024.
Aram-Alexandre Pooladian, Heli Ben-Hamu, Carles Domingo-Enrich, Brandon Amos, Yaron Lip- man, and Ricky 'T. Q. Chen. Multisample flow matching: Straightening flows with minibatch couplings. In International Conference on Machine Learning (ICML), 2023.
Emmanuel Quemener and Marianne Corvellee. Sidus-the the solution for extreme deduplication of an operating system. Limux J., 2013(235), November 2013. ISSN 1075-3583.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedi- cal image segmentation. In Medical Image Computing and Computer-Assisted Intervention (MIC- CAI), pp. 234-241. Springer, 2015,
Ernest Ryu, Jialin Liu, Sicheng Wang, Xiaohan Chen, Zhangyang Wang, and Wotao Yin. Plug-and- play methods provably converge with properly trained denoisers. In International Conference on Machine Learning (ICML), 2019.
Jascha Sohi-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In Francis Bach and David Blei (eds.), Proceed- ings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pp. 2256-2265, 2015.
Hendrik Sommerhoff, Andreas Kolb, and Michael Moeller. Energy dissipation with plug-and-pla priors. In Advances in Neural Information Processing Systems (NeurIPS) Workshop, 2019.
Jiaming Song, Arash Vahdat, Morteza Mardani, and Jan Kautz. Pseudoinverse-guine guided diffusion models for inverse problems. In International Conference on Learning Representations (ICLR), 2023.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. Interna- tional Conference on Learning Representations (ICLR), 2021.
Hannes Stark, Bowen Jing, Chenyu Wang, Gabriele Corso, Bonnie Berger, Regina Barzilay, and Tommi Jaakkola. Dirichlet flow matching with applications to DNA sequence design. In Interna- tional Conference on Machine Learning (ICLR), 2024.
Yu Sun, Brendt Wohlberg, and Ulugbek S Kamilov. An online plug-and-plat algorithm for regular- ized image reconstruction. IEEE Transactions on Computational Imaging, 5(3):395-4P) 408, 2019.
Julian Tachella, Dongdong Chen, Samuel Hurault, Matthieu Terris, and Andrew Wang. DeepInverse: A deep learning framework for inverse problems in imaging, 2023.
Hong Ye Tan, Subhadip Mukherjee, Junqi Tang, and Carola-Bibiam. Schönlieb. Provably convergent plug-and and--and-play quasi-Newton methods. SIAM Journal on Imaging Sciences, 17(2):85-8119 2024.
Matthieu Terris, Audrey Repetti, Jean-Christorn Christophe Pesquet, and Yves Wiaux. Building firmly nonex- pansive convolutional neural networks. In IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP). 2020.
14

---

Published as a conference paper at ICLR 2025
Alexander Tong, Kilian Fatras, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector- Brooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. Transactions on Machine Learning Research, 2024.
Singanallur V. Venkatakrishnan, Charles A. Bouman, and Brendt Wohlberg. Plug-and-play priors for model based reconstruction. In IEEE Global Conference on Signal and Information Processing (GlobalSIP), pp.
2013.
Cédric Villani. Optimal Transport - Old and New, volume 338. Springer, 2008
Yinhuai Wang. Jiwen Yu, and Jian Zhang. Zero-shot image restoration using denoising diffusion null-space model. In The Eleventh International Conference Learning Representations, 2023.
URL httpttts////peeeeieeeee http://.ee////tooeeeee/eee id=miiieee (ateeemmRe@gMMTT
Xinyi Wei, Hans van Gorp, Lizeth Gonzalez- Gonzalez-Carabarin, Daniel Freedman, Yonina C. Eldar, and Ruud G. van Sloun. Deep unfolding with normalizing flow priors for inverse problems. IEEE Transactions on Signal Processing, 70:2962-2971, 2022.
Zihui Wu, Yu Sun, Yifan Chen, Bingliang Zhang, Yisong Yue, and Katherine Bouman. Principled probabilistic imaging using diffusion models as plug-and-play priors. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL /ottttpee///pnnnvieew net/ met/foruu ceet/forum?icm id- SarprummidxxxHHFFFFF 7VNV
VVVVVVVVVVVV
Xingyu Xu and Yuejie Chi. Provably robust score-based diffusion posterior sampling for plug-and- play image reconstruction. arXiv preprint arXiv:2403. 17042, 2024.
Ling Yang. Zixiang Zhang, Zhilong Zhang, Xingchao Liu, Minkai Xu, Wentao Zhang, Chenlin Meng, Stefano Emnon, and Bin Cui. Consistency flow matching: Defining straight flows with velocity consistency. arXiv preprint wxXvv2247.002398 2024.
Shuo Yang, Ping Luo, Chen-Change Loy, and Xiaoou Tang. From facial parts responses to face detection: A deep learning approach. In IEEE International Conference Computer Vision (ICCV), pp. 3676-3684 2015.
Bingliang Zhang, Wenda Chu, Julius Berner, Chenlin Meng, Anima Anandkumar, and Yang Song. Improving diffusion inverse problem solving with decoupled noise annealing. arXiv preprint arXiv:2407.01521, 2024a.
Kai Zhang, Wangmeng Zuo, Yunjin Chen, Deyu Meng, and Lei Zhang. Beyond a Gaussian denoiser: Residual learning of deep CNN for image denoising. IEEE Transactions on Image Processing, 26(7):3142-3155, 2017.
Kai Zhang, Yawei Li, Wangmeng Zuo, Lei Zhang, Lue Van Gool, and Radu Timofte. Plug-and-plaz image restoration with deep denoiser prior. IEEE Transactions on Pattern Analysis and Machine Intelligence, 44(10)6660-66376 2021.
Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang. The unreasonable effectiveness of deep features as perceptual metric. In IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 586-595, 2018
Yasi Zhang, Peiyu Yu, Yaxuan Zhu, Yingshan Chang, Feng Gao, Ying Nian Wu, and Oscar Leong. Flow priors for linear inverse problems via iterative corrupted trajectory matching, arXiv preprint arXiv:2405. 18816, 2024b.
Yuanzhi Zhu, Kai Zhang, Jingyun Liang, Jiezhang Cao, Bihan Wen, Radu Timofte, and Luc image restoration. In /EEE/CVF Con-
Van Gool. Denoising diffusion models for ference on Computer Vision and Pattern Recognition (CVPR) Workshop, pp. 1219-1229, 2023,
15

---

Published as a conference paper at ICLR 2025
A APPENDIX
A.1 PROOF OF PROPOSITION I
Proof. We have by the definition of the denoiser (6) (it was also observed in Lee et al. (2023)
XI)a D.(,) X1 [|-] = (1 - (x. [l) (X1 - Xo))
If we assume that the denoising loss is zero, this yields for all E [0, that vx(X,) X X almost surely. By continuity this also follows for 1. By the same arguments in (Liu al. 2023 Theorem 3.6 iii)+ iv)) we have that both & + Xt) and + X, are solution to the flow ODE initialized Xo, since
X, Xo + tox( = Xo + f(vv(X.)dde a.s.
where we used that e is constant with time. This implies that 0,X, = (x(),) almost surely. Together with the uniqueness of the ODE solution, it follows that ((,, Xo) X, almost surely.
On the other hand, if (f; v) is a straight-line Flow Matching pair connecting Xo and X1, then we obtain that
(X,) = (f Xx)) = Ot = 0,X, = X1 Xo a.s.
Therefore the denoising loss is zero.
A.2 PROOF OF PROPOSITION 4
We provide the proof for Proposition [ here below. Note that the assumption of noise boundedness can be easily satisfied by clipping the noise. Additionally, naturally ensured numerically since floating-point numbers have finite precision.
Proof. By the algorithm and definition of Dr,,- have
In+1 = De., (u,) Un + (1 tn) ee (un),
where together with the definition of %m=
Un = (1 - tn + tn (xn - %/ "F = tn + (1 tn) (En - t, (F(x.))
Hence we obtain
|| - xn = (1- (.) - In (F()) +0 (un)ll-
By assumption on F and u and since both (xn),, and the noise e are bounded, the expression in the norm is bounded as well, say by M > 0. Then, by assumption on fn. we conclude
Σ n=0
| x, = M Σ
(1- 4m) <c,
n=0
so that (In)n is a Cauchy sequence and converges.
A.3 NUMERICAL RESULTS USING LATENT DIRICHLET DISTRIBUTION
One of the main practical advantages of Flow Matching over diffusion is that one can choose latent distribution with not Gaussian. Motivated by DNA design and other discrete data, there has been growing interest in using Dirichlet latent distribution in Flow Matching (Stark et al.) 2024 Boll et al., 2024 Davis et al. 2024). Here, we want to show that applying PnP-Flow with a Flow Matching model trained on a Dirichlet latent distribution still yields good reconstructions. For this, we use the MNIST dataset (LeCun et al.) 1998). rescaling each image to lie on the simplex, and train Flow Matching model. We then apply this model to inpainting and super-resolution tasks and present the results. Importantly, the goal of this experiment not to achieve state-of-the-an performance, but illustrate that algorithm generalizes to different latent distributions.
16

---

Published as a conference paper at ICLR 2025
For this, train Dirichlet Flow Matching model for 200 epochs with standard OT Flow Match- ing loss and a Dirichlet distribution with parameters (1...., 1) E I888. This latent distribution is also used in Stark et al. (2024) and amounts to the uniform distribution the 784-dimensional simplex. Note that the interpolation step now involves Dirichlet noise instead of Gaussian, as this noise is drawn according to Po- We reconstruct our images with y 1 and 300 steps. Remarkably, almost no modifications for the algorithm are needed. In particular, the generated images z almost lie per- fectly on the simplex without any normalization. We show the images in Figure [ We compare to (Ben-Hamu et al., 2024), where we adapt the regularization in the algorithm so that the optimized variable lies on the simplex, i.e., we use the loss
L(z) = lyy - ((((1, + XII 21,5) 1)1?
ij
for a single latent image 2, where f(1, is realized doing 5 Euler steps using the mid point rule and is a regularization constant.
As one can sec, the PnP Flow outperforms D-Flow on all tasks. Note that the other Flow Match-
ing baselines (Pokle et al. 2024 Zhang al., 2024b) are not usable since their algorithm heavily
depends on Gaussian paths.
Clean
Degraded
D.Flow
PnP Flow
2
2
4
7
2
2
2
+
2
2
G
7
0
9
4
7 0
9
4
7 2
0
9
4
7 2
0
9
4
2 2
5
7
5
2 2
5
7 7
5
5
7 7
5
5
7 7
5
2
7
2
2
2
4
2
2
MSF: 0.0008
MSF: 0.0003
MSE: 0.00009
7
2
2
4
2
2
>
2
2
7
0
9
4
7
0
9
4
7 2
0
9
4
2
5
7
8
31
2
5
7
8
5
7
8
2
2
7
2
2
7
2
2
7
427
MSE: 0.0064
MSE: 0.0018
MSE 0.0012
7
2
2
4
>
1/
2_
2
3
4
2
2
2
7
0
9
4
G
H1
7
0
9
K
9
0 5 2
8
6
2
5
7
5
C
0
5
3
R
2
8
2
2
7
7
2
2
7
2
7
1
MSE: 0.0049
MSE: 0.0042
MSE: 0.0029
Figure 5: Dirichlet Flow Matching experiment Simplex-MNIS for denoising (Ist row), super- resolution (2nd row), box- box-inpainting (3rd row). We measure the reconstruction eтror as the mean L2 distance (called MSE) between ground truth and reconstruction averaged over the 16 images.
A.4 SENSITIVITY TO INITIALIZATION
Interestingly, our method is inherently independent of the algorithm's initialization. For the image restoration problems considered this work, initialization is not so much of a concern because reasonable starting point for the solution is given by H y, where y is the observation and H is the degradation operator. However, for more complex problems where y is far from resembling natural image, such in CT reconstruction (Guo et al. 2016) phase retrieval (Manekar al. 2020), this property becomes much more relevant.
In contrast, competing methods more sensitive to initialization and cannot be started from any value. As recommended in the paper, our implementation of OT-ODE is initialized with foy + (1 fo)e where A/ Ta) and is the initialization time. The latent variable in D-Flow is initialized
17

---

Published as a conference paper at ICLR 2025
aT (H' y) + (1 a)c, where T is the reverse flow, <~ NO Ia) is random Gaussian noise, and a E (0,1) is blending coefficient. flow-Priors, on the other hand, initialized with random noise and, like our method, does not depend "good" initialization.
In Figure 6 we illustrate the impact changing the standard initialization for all methods to black image Gaussian deblurring task, comparing the robustness of each approach to poor initialization.
Clean
OT-ODE
D-Flow
Flow-Priors
Pal-Fllw
Figure 6: Comparison of restoration methods on the CelebA dataset, for a Gaussian deblurring task and for two different initializations: default initialization recommended for each method (Ist row), initialization set to the zero image (2rd row).
A.S EXPERIMENTS WITH RECTIFIED FLOWS
The image restoration experiments we carry out the core of the paper use a pre-trained OT Flow Matching model. Yet, the theory behind method holds for any straight-line Flow Matching model. In this section, we show how our method PnP-Flow compares to Flow Priors and OT-ODE using. for all three methods, pre-trained Rectified Flow model Liu al. (2033 dataset (Karras et al. 2018) with image dimension 256 x 256. We use the checkpoint provided by the Rectified Flow repository We were not able
CelebA-HQ
D-Flow using the Rectified Flow model.
Quantitative results are in Table and qualitative results on paintbrush inpainting are displayed in Figure
Table 4: Comparisons of state-of-the-art methods on different inverse problems on the dataset CelebA-HQ. Results averaged across 100 test images.
Denoising
Deblurring
Super-res. 0.05, x4
Rand. inpaint. 0.01, 70%
Box inpaint. 0.05, 80 x
Method
0.2
σ 0.05, a 3.0
PSNR
SSIM
PSNR
SSIM
PSNR
SSIM
PSNR
SSIM
PSNR
SSIM
Degraded
20.00
0.277
23.94
0.514
9.98
0.040
12,28
0.197
22.14
0.726
OT-OδE
29.38
0.731
25.47
0.589
24.24
0.578
26.59
0.639
26.68
0.683
Flow- -ri-Prroos
25.73
0.524
28.09
0.804
23.64
0.494
33.26
0.935
29.89
0.817
InP-Flow (ours)
32.65
0.891
28.57
0.815
25.92
0.762
33.60
0.935
31.11
0.928
A.6 PROGRESSION OF THE PNP-FLOW RECONSTRUCTION WITH TIME
Figure presents the progression of the reconstruction outputed by the PnP-Flow with respect to time.
hutp://ghthh..cmmmmmoiaaallllllf
18

---

Published as a conference paper at ICLR 2025
Clean
Degraded
OT-ODE
Flow-Priors
PaP-Flow
OCIE oss
OCIE
OCIE
OCIE
OCIE
1s
115
ass
X9400
X9400
XBADO
X3400
XBADO
Vl
V
M
V
Vi
Figure 7: Comparison of restoration methods on the CelebA-HQ dataset for paintbrush inpainting using Rectified Flow model.
Noisy
t 0.0
0.1 f= 0.2
t= 0.3
U.4
PSNR: 11.23
8.23
13.98
18.01
21.02
23.13
== 0.5
f 0.6
t 0.7
t= 0.8
f 0.9
t 1.0
25,80
27.35
29.23
31.42
34.40
36.78
Figure 8: Results for random inpainting using PnP-Flow across different iterations (time steps) with corresponding PSNR values. As expected. in the early iterations, the output resembles natural face but diverges from the noisy observation. PSNR improves progressively with each iteration.
A.7 ADDITIONAL VISUAL RESULTS
Here, we provide additional visual results, where we take the same ground truth image for the different inverse problems, see Fig. and Fig. 10 for CelebA dataset and Fig. 11 and Fig. 12 for AFHQ-cat dataset.
19

---

Published as a conference paper at ICLR 2025
Clean
Degraded
PnP-Diff
PnP-GS
OT-ODR
D-Flow
Flow-Priom
PaP-Flow
P05k 2002
FSNR 1004
PSNN: 22.11
PONK 1929
PONK
FS5E 20.4 27.45
PANIE
PSNIL
PSNS 2013
PSNE
PSNI 2031
PSNK J0 IV
PSNIE
N/A
N/A
Figure 9: Comparison of image restoration methods on CelebA: denoising (Ist row), Gaussian deblurring (2nd row), super-resolution (3rd row), free-form inpainting (4th row).
Clean
Degraded
PnP-Diff
PnP-GS
OT-ODE
D-Flow
Hoo-Prioos
PnP-Flow
Fank
PING 5030
Paae 1500
10NN 2010
73NAL 2011
FS5N 20.34
PONTL 24.55
7S5N4 20100
PONNX 20.0
PANK
N/A
N/A
Figure 10: Comparison of image restoration methods on CelebA: denoising (Ist row), Gaussian deblurring (2nd row), super-resolution (3rd row), free-form inpainting (4th row).
20

---

Published as a conference paper at ICLR 2025
Clean
Degraded
PnP--Diff
PnP-GS
OT-ODE
D-Flow
How-Prinrs
PnP-Flow
Pshile
P5Sile 1331
75500 13:93
FSS0 2035
FANK 15:31
FSSE 72.27
FS5R 10335
PANE: 31
FSNE 20.33
Fabxx 21:35
N/A
NA
FANK 2112
FENE 20012
PSNN 20.6
FSNE 2436
Pahox 21J11
Figure 1I: Comparison of restoration methods on AFHQ-Cat: denoising (Ist row), Gaussian de- blurring (2nd row), super-resolution (3rd row), free-form inpainting (4th row).
Clean
Degraded
PoP-Diff
PnP-GS
OT-ODE
D-Flow
Flow-Priors
PnP-Flow
PANE
PSNIC
ISNI 29.40
FSNIC 25
RINE
PANS
N/A
PANKL 1524
PANK 2514
PANK 2015
FANIL 25.1
PAND 2002
Figure 12: Comparison of image restoration methods on AFHQ-Car denoising (Ist row), Gaussian deblurring (2nd row), super-resolution (3rd row), free-form inpainting (4th row).
21

---

Published as a conference paper at ICLR 2025
A.8 OPTIMAL PETER-PARAMEEERS VALUES
We implemented the Gradient Step denoiser from Hurault et al. (2022a b) following the repository For random inpainting, the PnP algorithm is the Half Quadratic Splitting with the parameters given in the paper (Hurault et al. 2022ap. For denoising, we only apply the trained denoiser once, feeding with the true noise level σ. For all other methods, we use Proximal Gradient Descent (PGD), as prescribed in Hurault et al. (022b). Note that we do not constrain the Lipschitz constant of the denoiser. For the PGD ease, we tune three hyper-panee parameters: the learning rate 7 E (0.99, 2.0) in the gradient step (which is related to the regularization parameter À in the paper), the inertia parameter a (0.3, 0.5, 0.8, 1.0] which coresponds to the relaxed denoiser De = DD + (1 - a)Id, and the factor of E {1., 1.2, 1.5, 1.8, 2., 4. J.J so that the denoiser gets as noise map of (where is the true noise level). We also considered the number of iterations as a hyperparameter that we tuned on the validation set (with maximum number of iterations fixed to 100).
F
For PnP-Diff, we tuned two parameters: the regularization parameter À e {1.0, 10.0, 100.0, 1000.0) and the blending parameter Ç {0.1, 0.3, 1.0, The num- ber of iterations was fixed to 100.
For D-Flow, we adjusted the blending parameter for initialization E 0.3,, 0.5, and the regu- larization parameter E {0.1, 0.01, 0.0013. We observed that the PSNR of the reconstruction did not consistently increase across iterations (the method does not always converge); thus, fine- tuned the value of the last iteration while keeping it below 20 for computational efficiency. The number of iterations for the inner LBFGS optimization was set to 20, and the number of Euler step when solving the ODE to 5, as recommended in the paper
et al. (2024).
For OT-ODE, tuned the initial time to {0.1, 0.2, 0.3, 0.43 and the type learning step y (either Vt constant). The number of iterations was 100.
For Flow-Priors, we considered -as described in the paper- the two {10-3, 10 2, 10 and E {102, 10a, 10a, 10+} which respectively correspond the step size for the gradient descent and the guidance weight (weight put on the data likelihood). The number of iterations was set 100 and the number of iner iterations K =
E
Finally, for our method PnP-Flow, we adjusted the exponent in the learning rate a E (0.01, 0.3, 0.5, 0.8, 1.0) and the number time steps N E {100, 200, 500). When increasing N beyond 100 resulted in less than 0.2 dB improvement in PSNR, we set N = 100 for computa- tional efficiency. We average the output of the denoising step in Algorithm 3 over realizations of the interpolation step. To speed up the algorithm, this number can be reduced to I with only minor impact on performance.
A.9 VISUALIZATION OF EACH STEP OF THE ALGORITHM
In Figure 13, we show the output of the PnP-Flow algorithm at different stages for an image from the CelebA dataset: after the gradient step, the interpolation step, and the denoising step.
A. ADDITIONAL. EVALUATION METRIC
Here, we add a comparison of the models in the LPIPS metric Zhang et al. (018), which measures the feature distance between the ground truth and generated images via some pretrained classifica- tion network. The results are contained in Table Table g None of the methods stands out as a clear winner with respect to the LPIPS metric. However, PnP- shows good performance, often being the second or third best method. In Blau & Michaeli (2018). it is argued that perceptual and distortion seem be odds with one another: indeed, one can show that the optimal distortion performance (optimizing the PNSR) does not yield optimal perceptual performance (more in line with LPIPS).
22

---

Published as a conference paper at ICLR 2025
Table 5: Hlyper-parameter used for all methods on the CelebA dataset. based on the highest PSNR on the validation split.
The values were selected
Denoising
Deblurring
Super-res.
Rand. inpaint.
Box inpaint.
PnP-Diff
Ç (blending)
1.0
1.0
1.0
1.0
N/A
À (regularization)
1.0
1000.0
100.0
1.0
N/A
PnP-GS
(learning rate)
2.0
2.0
N/A
(inertia param.)
1.0
0.3
1.0
N/A
oy (factor for noise input)
1.0
1.8
3.0
N/A
Fiter (number iter.)
20
N/A
OT-ODE
(initial time)
0.1
0.1
vf
vt
constant
constant
va
À (regularization)
100
1,000
10,000 0.1
10,000 0.01
10,000
7 (learning rate)
0.01
0.01
0.01
D-Flow
À (regularization)
0.001
0.001
0,001
0.01
0.001
(blending)
0.1
0.1
0.1
0.1
0.1
Tlier (number iter.)
3
10
PnP-Flow
a (learning rate factor)
0.8
0.01
0.3
0.01
0.5
N (Number of time steps)
100
100
100
100
100
Table
used for all methods on the AFHQ-Cat dataset. The values were selected based on the highest PSNR on the validation split.
Denoising
Deblurring
Super-res.
Rand. inpaint.
Box inpaint.
PnP-Diff
C (blending)
1.0
1.0
1.0
1.0
N/A
À (regularization)
1.0
1000.0
100.0
1.0
N/A
PnP-GS
Y (learning rate)
2.0
2.0
N/A
(inertia param.)
1.0
0.3
1.0
N/A
ay (factor for noise input)
1.0
1.8
5.0
N/A
Tiieer (number of iter.)
1
60
50
23.
N/A
OT-ODE
to (initial time)
0.3
0.3
0.1
0.1
0.1
Y
v7
v7
constant
constant
vī
Flow-Prion
(regularization)
100
1,000
10,000
10.000
10,000
(learning rate)
0.01
0.01
0.1
0.01
0.01
D-Flow
À (regularization)
0.001
0.01
0.001
0.001
0.01
a (blending)
0.1
Flier (number of iter.)
3
20
20
20
9
PnP-Flow
(leariing rate factor)
0.8
0.01
0.01
0.01
0.5
(Number of time steps)
100
500
500
200
100
23

---

Published as a conference paper at ICLR 2025
Noisy
Reration
Likellhood step: enforce
Interpolation step: reproject the tteraee onto aspropriste space expecled by of
Denaising step: push towards data distribution
Iteration 20
tteaaion 80
output
Figure 13: Progression of the output of the algorithm at intermediate steps (log-likelihood step, reprojection onto the path, and denoising) for random inpainting task.
Table 7: Comparisons of state-of-the-art methods on different inverse problems on the dataset CelebA using the LPIPS perceptual metric Zhang al. (2018). Results averaged across 100 test images.
Denoising 0.2
Deblurring 0.05, a
Super-res. 0.05, x4
Rand. inpaint. 0.01, 70%
Box inpaint.
Method
0.05,
Degraded
0.373
0.125
0.827
1.034
0.213
PnP-Diff
0.060
0.060
0.034
0.025
N/A
PnP-GS
0.035
0.040
0.053
0.085
N/A
ÖT-ÖDE
0.110
0.096
0.091
0.137
0.102
D-Flow
0.078
0.053
0.027
0.021
0.035
Flow-Prioe Pow-Priors
0.137
0.056
0.101
0.018
0.048
PnP-Flow (ours)
0.056
0.046
0.055
0.021
0.043
A.11 HANDLING NON- GANSSSAN NOISE MODELS
Our approach is capable of handling various types of degradation noise, beyond Gaussian. gen- eral, any noise model corresponding to a differentiable likelihood be handled by method effortlessly. Note that this is not the case of all D-Flow and OT-ODE may be possible to adapt these methods to other kinds of noises, doing so would require additional modifications, which we leave for future work.
methods. In particular, the methods designed under the assumption of standard Gaussian noise. Although it
In the following experiment, we consider a super-resolution task with a degradation model incorpo- rating Laplace noise:
y Lapplace (Hx),
24

---

Published as a conference paper at ICLR 2025
Table 8: Comparisons of state-of-the-art methods on different inverse problems on the dataset AFHQ-Cat using the LPIPS perceptual metric Zhang et al. (2018). Results are averaged across 100 test images.
Denoising 0.2
Deblurring
Super-res.
Rand. inpaint. 0.01, 70%
Box inpaint.
Method
0.05,
3.0
0.05, x4
0.05, 80 80
Degraded
0.524
0.460
0.888
1.069
0.206
PnP-Diff
0.194
0.354
0,438
0.064
N/A
PnP-GS
0.075
0.360
0.384
0.116
OT-ODE
0.190
0.226
0.193
0.226
0.187
D-Flow 0.177
0.178
0.207
0.049
0.091
Flow-Priors 0.159
0.191
0.279
0.047
0.122
PnP-Flow (ours)
0.170
0.341
0.258
0.043
0.118
where H : _ + IR. is 2-downsampling linear operator. Laplace noise, which has heavier tails than Gaussian noise, is well-suited for modeling data with outliers. In this context, the data-fidelity function is defined as (z) = Hr yll and differentiable almost everywhere.
For each method, tune the saper-parameters via grid search, selecting the configuration that yields the highest PSNR, similar to our previous experiments. We report the results in terms of PSNR, SSIM, and LPIPS in Table g We also show the visual results in Figure 14
Table 9: Comparisons of state-of-the-art methods on Laplace super-resolution experiments on the dataset CelebA. Results are averaged across 100 test images.
Laplace super-resolution
Laplace super-resolution
Method
0.1
0.3
PSNR
SSIM
LPIPS
PSNR
SSIM
LPIPS
Degraded
8.82
0.086
0.845
8.76
0.040
0.865
PnP-Diff
29.67
0.866
0.059
24.82
0.730
0.268
PnP-GS
28.50
0.828
0.079
25.31
0.725
0.185
Flow-Priors
27.53
0.721
0.102
21.72
0.454
0.376
PnP-Flow (ours)
30.30
0.895
0.063
26.50
0.809
0.108
Clean
Degraded
PaP-Ditt
PaP-GS
Flow-Prine
PaP-Flow
Figure 14: Example of reconstructions obtained on super-resolution tasks with Laplace noise
(a = 0.3), on the CelebA dataset.
12 DISCUSSION OF RELATED DIFFUSION APPROACHES
In He et al. (2024), based on the seminal work Chung et al. 2003),, a similar algorithm proposed in the diffusion case (although for latent models): Our algorithm closely resembles their pixel space algorithm with a different noise update rule. However, note that they do not relate it to PnP methods, and that outlined several theoretical and practical advantages of flow matching over diffusion models. Further similar ideas are explored Wu et al. (2024): [Zhang et al.] (2024a); Xu & Chi 2024): Li et al. (024), which focus posterior sampling using diffusion models. These works
25

---

Published as a conference paper at ICLR 2025
employ a similar decoupling of the likelihood (data- (data-consistency) sampling step and a denoising diffusion sampler (referred to as the decoder in [i et al. (2024)) that acts as an implicit image prior. Additionallyy, Wang et al. (2023) introduces particularly efficient null-space approach. Instead of making steps on the log likelihood like Chung et al. (2023), they optimize in the null space of the
linear operator directly. Such ideas could in principle also be applied to algorithm, which leave for future work.
A.13 ADDITIONAL EXPERIMENTS WITH DIFFUSION-BAREE METHODS
Table 10: Comparison of PnP-Flow with diffusion- based methods various inverse problems using the CelebA dataset. Results are averaged over 100 test images.
Denoising
Deblurring 0.05,
Super-resolution 0.06,
Random Inpainting 0.01, 70% mask
Method
0.2
1.0
PSNR
SSIM
LPIPS
PSNR
SSIM
LPIPS
PSNR
SSIM
LPIPS
PSNR
SSIM
LPIPS
Degraded
20.00
0.348
0.373
27.67
0.740
0.125
7.53
0.012
0.827
11.82
0.197
1.034
PnP-Diff
31.00
0.883
0.060
32,45
0.910
0.060
51.20
0.353
0.034
33.45
0.917
"0.025"
DDRM
31.46
0.895
0.071
31.44
0.901
0.080
31.56
0,909
0.046
30.31
0,902
0,040
DPS
31.52
0.887
0.068
31.42
0.917
0.038
29.18
0.841
0,588
26.83
0.858
0.122
PriPHoow (Öurs)
32.45
0.911
0.056
34.51
0.940
0.046
31.49
0.367
0.055
33.54
0.953
"0.021
Table 11: Comparison of PnP-Flow with diffusion- -based methods on various inverse problems using the AFHQ- dataset. Results are averaged over 100 test images.
Denoising
Debluriing
Super-resolation = 0.05, x1
Random Inpainting
Method
===22 0.2
0.05, = 3.0
= 0.01, 70% mask
PSNR
SSIM
LPIPS
PSNR
SSIM
LPIPS
PSNR
SSIM
LPIPS
PSNR
SSIM
LPIPS
Degruded
20.00
0,319
0,524
23.77
0,514
0.460
10.74
0.042
0.888
13.35
0.234
1.069
PuP-Diff
30.27
0.835
0.194"
27.97
0.764
0.354
23.22
0.601
0.438
31.08
0.882
00.044"
DDRM
30.60
0.838
0.240
25.89
0.680
0.430
27.50
0.746
0.329
30.24
0.857
0.109
DPS
27.64
0,743
0.327
24.31
0.673
0.390
27.83
0.755
0.322
27.23
0.739
0.334
Pnow (Gurs)
31.65
0.876
0.170
21.62
0,763
0.341
26.75
0.774
0.258
"32.98
0.930
"0,04
In Table 10 and Table IT. we provide extensive benchmarks using diffusion-banae methods (PnP- Diff Zhu et al. (2023). DDRM Kawar et al. (2022) and DPS Chung et al. (2023) ), respectively
on the CelebA and AFHQ-Cat datasets. Qualitative results are given in Figure [5] All diffusion-
based methods use the same pre-trained model from Choi et al. (202I), implemented in the DeepInv library Abbraraa Tachelll et al. 2023). Note that the pre-trained diffusion model was trained on the FFHQ dataset KKrrra al. 2019), making the comparison indirect. For all methods, gridsearch over hyperparameters has been conducted (same methodology as followed for the other experiments, see Appendix (A.).. More precisely, for DDRM, we adjusted the value of 17 and 1/8 with grid search over {0.7, 0.9, 1.0] following the ablation study in Kawar et al. (2022). For DPS, we adjusted the value of 7) E {0.5, 0.7, 0.8, 0.85, 0.9, 0.95) which controls the stochasticity.
A.14 EXTENSION TO BLIND INVERSE PROBLEMS
In this Section, we describe how our algorithm can be used for blind deconvolution Levin et al. (209). i.e., in deblurring problem, where the forward operator is unknown. Here, we assume that we know the kernel size and that the kernel is the same across channels. The true forward operator is given by a Gaussian blur and we attempt to simultaneously reconstruct blur and "true" image. This we do as follows: We loop over discrete time sequence and after each time-step in Algorithm
3 we run a gradient descent update of the "Icarned" kernel on the log likelihood loss lly (earned)))) [12 2 with the Adam optimizer (Kingma & Ba] 2017). The result be seen for one examplary image in Figure 16 where we see that our reconstruction is clearly improving upon the blurry measurement. In Figure 17 we repeat this experiment for a random inpainting task, learning
26

---

Published as a conference paper at ICLR 2025
Clean
Degraded
Pii--Dif
DDRM
PaP-Flow
Figure 15: Up: example of reconstructions obtained on a Gaussian denoising task (a = 0.2) on the CelebA dataset. Bottom: example of reconstructions obtained random inpainting task (70% of masked pixels) on the AFHQ-Cat dataset.
both the reconstruction and the parameterized mask. This seems to be an interesting avenue for future research.
Clean ground truth.
Blurry measurement.
Reconstruction with our method.
Figure 16: Application of our method to a blind deconvolution problem.
Clean ground truth.
Blurry measurement.
Reconstruction with our method.
Figure 17: Application of our method to blind masking problem.
27


Reviews:
[{"review_id": "IFPB0tusgK", "reviewer": "Reviewer_pWBL", "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.", "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.\n\n2. 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.\n\n3. 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.\n\n4. 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. \n\n5. For real-world data with unknown degradation, it would be important to understand how well this method generalizes. \n\n6. It would strengthen the paper if the authors included examples of failure cases. \n\n7. 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.", "limitations": "", "rating": 3, "confidence": 5, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "1. Low memory usage, making it suitable for high-resolution images.\n2. 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.\n2. The details of the proposed method are insufficient.\n3. The experiment section should be improved. Please refer to the details below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "zg3oIm30kI", "reviewer": "Reviewer_n2r9", "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.", "questions": "1. Could you add comparisons with [1] and [2], or explain why those comparison are missing?\n2. Could you comment on the potential applicability/extension of your method to the blind case?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The method is training-free which makes it computationally practical.\n2. 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].\n2. The proposed method is non-blind (assume the full knowledge of the degradation model) which limits its applicability.\n\n\n[1] Wang et al. Zero-Shot Image Restoration Using Denoising Diffusion Null-Space Model. ICLR 2023\n\n[2] Chung et al. Diffusion Posterior Sampling for General Noisy Inverse Problems. ICLR 2023", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "B9UNg1St4B", "reviewer": "Reviewer_fVXi", "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.", "questions": "See weakness", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1, It proposes a design a time-dependent denoiser based on a pre-trained velocity field v learned through Flow Matching\n\n2, 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\n\n3, 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? \n\n2, 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)?\n\n3, Can you visualize all the intermidate resutls of all three steps for all time steps? It could better help readers understand the method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "e2bcen88PY", "reviewer": "Reviewer_VJZf", "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.", "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?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. A new plug-and-play method based on flow matching is proposed in this paper.\n\n2. The paper is well-written.\n\n3. The derivations in this paper are rigorous.\n\n4. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)