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

• Image restoration seeks an unknown image from a degraded observation, and the problem is usually ill-posed and high-dimensional.The observation is modeled using a degradation operator and additive noise.Source: §1 · p. 1
• PnP methods use denoisers within optimization schemes, while generative models learn image priors directly from data.The paper notes that recent generative models offer a more sophisticated framework than hand-crafted or neural denoisers.Source: §1 · pp. 1–2
• Earlier invertible-transport approaches use explicit priors, but Real-NVP yields poor results on large images; diffusion and Flow Matching models scale better.These latter models establish a path between latent and target distributions that can be optimized in a simulation-free way.Source: §1 · p. 2
• The method integrates an implicit Flow Matching prior into a custom denoiser plugged into a Forward-Backward algorithm, avoiding backpropagation through the ODE.Direct MAP optimization with the change-of-variables formula faces numerical challenges due to ODE backpropagation.Source: §1 · p. 2
• Across two datasets and four tasks, PnP-Flow consistently outperforms state-of-the-art Flow Matching-based and PnP methods on PSNR and SSIM.The tasks are denoising, deblurring, inpainting, and super-resolution.Source: §1 · p. 2

2.1 PLUG AND PLAY

• PnP methods replace the regularization proximal step with an off-the-shelf denoiser within proximal-splitting algorithms such as Forward-Backward Splitting.For FBS, convergence to a minimizer is guaranteed under the stated differentiability and stepsize conditions.Source: §2.1 · pp. 2–3
• Many neural-denoiser PnP-FBS algorithms do not converge, while convergence guarantees usually assume non-expansiveness, which can harm network expressiveness when enforced.The cited constraints include weight clipping, spectral normalization, and averaged-operator constructions.Source: §2.1 · p. 3

2.2 FLOW MATCHING

• Flow Matching defines a probability path from latent to data distributions and learns a velocity field whose ODE transports samples along that path.For straight-line flows, the path interpolates between paired latent and data samples; OT couplings are a particular case.Source: §2.2 · pp. 3–4
• Conditional Flow Matching training avoids needing the true velocity field and requires sampling from the chosen coupling, which may be independent or approximately optimal.The paper states that the conditional loss equals the Flow Matching loss up to a constant.Source: §2.2 · p. 4

3.1 DENOISING OPERATOR FROM FLOW MATCHING

• With an ideal velocity field, the time-dependent denoiser estimates the conditional mean of the target sample given a point on the flow path; its denoising loss is zero exactly for straight-line flows.The operator minimizes the squared-error objective for predicting the target sample from the interpolated point.Source: §3.1 · p. 5
• The denoiser is designed for straight paths, whereas diffusion models generally induce non-straight paths that do not match the desired linear interpolation.The paper notes that some non-straight Flow Matching paths can be rectified to straighter ones.Source: §3.1 · p. 5

3.2 PNP FLOW MATCHING ALGORITHM

• Each iteration takes a data-fidelity gradient step, interpolates the result with noise sampled from P0, then applies the time-dependent denoiser.The interpolation reprojects the gradient-updated result onto the flow path; the noise sample is not coupled to that result via π.Source: §3.2 · pp. 5–6
• The learning rate decreases with time as γ_t = (1 − t)^α to balance data fit and denoising; averaging several noise realizations slightly improves numerical results.The averaging makes the algorithm output deterministic.Source: §3.2 · p. 6
• If the iterates are bounded and the stated assumptions hold, the algorithm converges when the time sequence satisfies the specified summability condition and γ_n = 1 − t_n.The proposition assumes continuously differentiable F, continuous learned velocity field, and a latent distribution with bounded support.Source: §3.2 · p. 6

Pre-trained Flow Matching methods

• OT-ODE [Pokle et al., 2024] derives a conditional velocity field with Tweedie’s formula and samples posterior images by solving an ODE.It assumes a Gaussian latent distribution and does not require retraining.Source: §4 · p. 6
• Flow-Priors [Zhang et al., 2024b] solves time-dependent MAP subproblems using a closed-form gradient, but assumes independent coupling and a Gaussian latent and computes an expensive trace.Source: §4 · p. 6
• D-Flow [Ben-Hamu et al., 2024] optimizes a latent loss through the flow ODE; its non-convexity can impede finding a global minimizer, and backpropagation is computationally expensive.The latent and original problems are theoretically equivalent because the flow map is invertible.Source: §4 · pp. 6–7

(PnP) Diffusion methods

• PnP diffusion methods replace a regularization proximal step with a pre-trained diffusion denoiser, with related variants adding interpolation or a gradient data-fidelity step.The paper presents PnP-Flow as the first PnP method based on Flow Matching.Source: §4 · p. 7

5.1 BASELINES

• Baselines include three Flow Matching restoration methods, PnP-Diff, and PnP-FBS; the authors implemented the Flow Matching baselines from their published descriptions.Source: §5.1 · p. 7

5.2 EXPERIMENTAL SETUP

• Experiments use CelebA and AFHQ-Cat, with 100 test images per task across denoising, deblurring, super-resolution, box inpainting, and random-pixel inpainting.The study trains Flow Matching models per dataset and tunes hyperparameters on validation PSNR; its denoising averages five interpolation realizations.Source: §5.2 · pp. 7–8, §5.3 · p. 8

5.3 MAIN RESULTS

• Other methods show task-dependent strengths: D-Flow and Flow-Priors do well in box inpainting but struggle with denoising, while PnP-GS, PnP-Diff, and OT-ODE show the reverse pattern.Source: §5.3 · p. 8
• Visually, PnP-Flow produces realistic, artifact-free images that can be slightly over-smoothed; D-Flow sometimes hallucinates, Flow-Priors adds artifacts, and OT-ODE struggles with image generation.Source: §5.3 · p. 8, Figure 3 · p. 9, Figure 4 · p. 9
• PnP-Flow ranks first or second in both PSNR and SSIM across all tasks and datasets, and the tables report results for CelebA and AFHQ-Cat.Source: Table 1 · p. 8, Table 2 · p. 8, §5.3 · p. 8

5.4 PRACTICAL ASPECTS

• Starting at t0 = 0 makes the initial interpolation pure noise, and the authors report that PnP-Flow’s performance is independent of initialization.Source: §5.4 · p. 10
• PnP-Flow can use any latent distribution, unlike the other Flow Matching restoration methods described by the authors.They illustrate a non-Gaussian case with a Dirichlet latent.Source: §5.4 · p. 10
• Although its denoiser is rooted in straight-line flows, the method also works with pre-trained Rectified Flows, with similar performance reported to Section 5.3.Source: §5.4 · p. 10

6 CONCLUSION

• PnP-Flow Matching uses few hyperparameters and minimal memory and computational resources while performing well across various inverse problems.It supports different latent distributions and flexible initialization.Source: §6 · p. 10
• Its reconstructions tend to be smooth because the denoiser is a mean square estimator.The authors describe this as a common tradeoff in inverse-problem regularization.Source: §6 · p. 10
• Future directions include other measurement-noise models, blind inverse problems, categorical latent distributions, and posterior sampling.The conclusion notes preliminary Laplace-noise results and promising blind-inverse-problem results in the appendices.Source: §6 · p. 10

A.1 PROOF OF PROPOSITION 1

• Assuming zero denoising loss, the velocity along the path equals X1 − X0 almost surely for every t in [0, 1].Continuity extends the result to t = 1.Source: §A.1 · p. 17
• ODE uniqueness implies that the learned flow f(t, X0) equals Xt almost surely.Both are solutions initialized at X0.Source: §A.1 · p. 17
• For a straight-line Flow Matching pair, the velocity along the flow trajectory equals the endpoint displacement X1 − X0 almost surely.The proof identifies the velocity with the time derivative of the straight-line path.Source: §A.1 · p. 17
• The proof concludes that the denoising loss is zero for a straight-line Flow Matching pair.Source: §A.1 · p. 17

A.2 PROOF OF PROPOSITION 4

• Bounded noise can be ensured by clipping and is also naturally enforced numerically by finite floating-point precision.Source: §A.2 · p. 17
• The proof rewrites the algorithm's intermediate iterate as a mixture of the current iterate and a noise-adjusted gradient step.It writes un = tn xn + (1 − tn)(εn − tn∇F(xn)).Source: §A.2 · p. 17
• Boundedness of the iterates and noise, together with assumptions on F and v, bounds the expression in the convergence estimate.The proof denotes this bound by M.Source: §A.2 · p. 17
• Under the stated assumptions, the iterates form a Cauchy sequence and converge.Source: §A.2 · p. 17

A.3 NUMERICAL RESULTS USING A LATENT DIRICHLET DISTRIBUTION

• The experiment tests generalization to a different latent distribution, not state-of-the-art performance, using MNIST images rescaled onto the simplex.Source: §A.3 · p. 17
• The model uses a uniform Dirichlet latent distribution on the 784-dimensional simplex, with Dirichlet noise in the interpolation step.It is trained for 200 epochs with a standard OT Flow Matching loss; reconstruction uses γ = 1 and 300 steps.Source: §A.3 · p. 18
• Generated images almost lie on the simplex without normalization, and the algorithm needs almost no modifications for this latent distribution.Source: §A.3 · p. 18
• Other Flow Matching baselines cited in the experiment are not usable because their algorithms heavily depend on Gaussian paths.Source: §A.3 · p. 18
• PnP-Flow outperforms D-Flow on denoising, super-resolution, and box-inpainting.Figure 5 reports reconstruction error as mean L2 distance, called MSE, averaged over 16 images.Source: Figure 5 · p. 18, §A.3 · p. 18

A.4 SENSITIVITY TO INITIALIZATION

• PnP-Flow is inherently independent of initialization, unlike competing methods that are more sensitive and cannot start from any value.The comparison includes changing methods’ standard initializations to a black image for Gaussian deblurring.Source: §A.4 · pp. 18–19, Figure 6 · p. 19
• For the restoration problems considered, H⊤y usually gives a reasonable starting estimate, so initialization is less of a concern.The paper identifies CT reconstruction and phase retrieval as settings where H⊤y may be far from a natural image and initialization independence matters more.Source: §A.4 · p. 18
• Figure 6 compares Gaussian deblurring reconstructions using each method’s recommended initialization and a zero-image initialization.The comparison is framed as a test of robustness to poor initialization; the supplied passages do not state its outcome.Source: §A.4 · p. 19, Figure 6 · p. 19

A.5 EXPERIMENTS WITH RECTIFIED FLOWS

• On CelebA-HQ at 256 × 256, PnP-Flow, Flow Priors, and OT-ODE use a pre-trained Rectified Flow model; D-Flow was not run with it.Table 4 compares state-of-the-art methods across inverse problems, with results averaged over 100 test images.Source: §A.5 · p. 19, Table 4 · p. 19

A.6 PROGRESSION OF THE PNP-FLOW RECONSTRUCTION WITH TIME; A.9 VISUALIZATION OF EACH STEP OF THE ALGORITHM

• Early random-inpainting iterations resemble a natural face but diverge from the noisy observation; PSNR improves progressively, while Figure 13 traces gradient, interpolation, and denoising stages.The progression is shown over time steps and at intermediate algorithm stages.Source: Figure 8 · p. 20, §A.9 · p. 23, Figure 13 · p. 25

A.7 ADDITIONAL VISUAL RESULTS

• Additional qualitative panels reuse ground-truth images across CelebA and AFHQ-Cat restoration tasks, including denoising, deblurring, super-resolution, and free-form inpainting.The figures show task rows for both datasets.Source: §A.7 · p. 20, Figure 9 · p. 21, Figure 10 · p. 21

A.8 OPTIMAL HYPER-PARAMETERS VALUES

• For PnP-Flow, when increasing N beyond 100 improved PSNR by less than 0.2 dB, the authors set N = 100 for efficiency; averaging five interpolation realizations could be reduced to one with minor impact.The learning-rate exponent and number of time steps were adjusted.Source: §A.8 · p. 23

A.10 ADDITIONAL EVALUATION METRIC

• LPIPS comparisons show no clear winner; PnP-Flow is often second or third best, while perceptual quality and distortion can be at odds.Tables 7 and 8 report LPIPS comparisons for CelebA and AFHQ-Cat, respectively.Source: §A.10 · p. 23, Table 7 · p. 25, Table 8 · p. 26

A.11 HANDLING NON-GAUSSIAN NOISE MODELS

• The method handles differentiable likelihoods for non-Gaussian noise; unlike D-Flow and OT-ODE, it does not require additional modifications for such noise models.A Laplace-noise super-resolution example uses an ℓ1 data-fidelity that is differentiable almost everywhere.Source: §A.11 · pp. 25–26

A.12 DISCUSSION OF RELATED DIFFUSION APPROACHES

• The paper relates its likelihood-step and denoising-step design to diffusion approaches, and notes that null-space optimization could also be applied to its algorithm as future work.The cited diffusion methods decouple data consistency from a denoising sampler; the null-space approach is attributed to Wang et al. (2023).Source: §A.12 · pp. 26–27

A.13 ADDITIONAL EXPERIMENTS WITH DIFFUSION-BASED METHODS

• Diffusion-based methods are benchmarked on CelebA and AFHQ-Cat with 100-image averages, but the comparison is indirect because the shared pre-trained diffusion model was trained on FFHQ.Qualitative examples include Gaussian denoising and random inpainting.Source: §A.13 · p. 27, Table 10 · p. 27, Table 11 · p. 27