PnP-Flow: Unified Inverse Imaging Framework

• PnP-Flow is a unified framework that combines iterative Plug-and-Play methods with generative flow matching models to address inverse imaging tasks like denoising, deblurring, super-resolution, and inpainting.
• The algorithm alternates between data-fidelity gradient descent and a time-dependent denoising step, ensuring rapid convergence and memory efficiency.
• Empirical evaluations on benchmarks demonstrate that PnP-Flow achieves high restoration quality with lower computational cost compared to other flow-based or diffusion methods.

Plug-and-Play Flow Matching (PnP-Flow) is a unified algorithmic framework for imaging inverse problems that fuses classical Plug-and-Play (PnP) iterative methods with the generative capacity of pre-trained Flow Matching (FM) models. This allows leveraging expressive learned generative priors in a memory-efficient, modular manner, facilitating superior image restoration across a range of tasks, including denoising, deblurring, super-resolution, and image inpainting. PnP-Flow addresses both the limitations of classical PnP, which struggles with highly nontrivial or “generative” inverse tasks such as inpainting, and the inefficiencies of directly incorporating flow-based generative models into optimization pipelines (Martin et al., 3 Oct 2024).

1. Mathematical Formulation and Theoretical Foundations

PnP-Flow arises from the variational formulation of the imaging inverse problem: $\hat x = \arg\min_{x \in \mathbb{R}^d} F(x) + R(x)$ where $F(x)$ is the data-fidelity term (e.g., $F(x)=\frac12\|Hx-y\|^2$ for Gaussian noise), and $R(x)$ represents a regularization penalty approximating the negative log-prior. Classical PnP methods replace the proximal operator of $R(x)$ with a learned denoiser $D$, leveraging the power of deep neural networks to encode prior knowledge.

Flow Matching models define a time-dependent velocity field $v_t^\theta$ trained to minimize

$$\mathcal{L}_{\rm CFM}(\theta) = \mathbb{E}_{t, (x_0, x_1)} \left\| v_t^\theta((1-t)x_0 + t x_1) - (x_1 - x_0)\right\|^2$$

resulting in an ODE

$$\frac{\partial f(t, x)}{\partial t} = v_t^\theta(f(t, x)), \quad f(0, x) = x$$

which transports samples from a known latent $P_0$ (e.g., Gaussian) to a data distribution $P_1$.

PnP-Flow synthesizes these frameworks by constructing a time-dependent denoiser

$D_t(x) = x + (1-t)v_t^\theta(x)$

which approximates the minimum mean-squared error (MMSE) solution at time $t$, with $v_t^\theta$ trained via straight-line FM. In the ideal infinite-data limit,

$$D_t(x) = \mathbb{E}[X_1 \mid X_t = x], \quad X_t = (1-t)X_0 + t X_1$$

so $D_t$ recovers the posterior mean of the clean image conditioned on a noisy intermediate.

The PnP-Flow iterative algorithm alternates:

1. Gradient descent on the data-fidelity,
2. Reprojection via stochastic interpolation onto the FM path,
3. Denoising using the time-dependent FM denoiser.

This yields rapid convergence to a fixed point, and the process adapts smoothly between data and prior depending on the temporal schedule.

2. Algorithmic Structure and Implementation

At each iteration $n$ (out of $N$), with time index $t_n = n/N$ and step size $\gamma_n = (1-t_n)^\alpha$, the core update rules are: $$\begin{aligned} z_n &= x_n - \gamma_n \nabla F(x_n) \ \tilde z_n &= (1 - t_n)\varepsilon_n + t_n z_n, \qquad \varepsilon_n \sim P_0 \ x_{n+1} &= D_{t_n}(\tilde z_n) = \tilde z_n + (1-t_n) v_{t_n}^\theta(\tilde z_n) \end{aligned}$$ Typical hyperparameters are $N = 100$–$200$, $\alpha \in (0, 1]$. Multiple noise draws for $\varepsilon_n$ provide enhanced robustness; no backpropagation through ODEs is required, and memory use is minimal (one pass through $v_t^\theta$ per iteration) (Martin et al., 3 Oct 2024).

In an advanced continuous-limit analysis, the method is modeled as a stochastic differential equation (SDE): $$dX_t = b_t(X_t) dt + \sqrt{2} dW_t, \quad b_t(x) = -x - \beta(t)\nabla F(x) + u_t(x)$$ where $u_t(x)$ is the FM velocity field and $\beta(t)$ is a time-dependent weighting. This view facilitates theoretical analysis and design of step-size schedules, Lipschitz regularization of $u_\theta$, and extrapolation-based acceleration (Jia et al., 3 Dec 2025).

3. Computational Efficiency and Scalability

PnP-Flow achieves high computational efficiency compared to alternative flow- or diffusion-based priors by:

• Completely avoiding backpropagation through the ODE solver (unlike D-Flow),
• Not computing Jacobian traces (avoiding heavy trace estimation as in Flow-Priors),
• Only requiring simple forward passes through the FM neural network and single-step gradient computations.

Empirical resource usage for deblurring ($128 \times 128$) on CelebA is summarized as follows:

Method	Runtime (s)	Memory (GB)
OT-ODE	1.5	0.65
Flow-Priors	16	3.0
D-Flow	32	6.0
PnP-Flow	3.4	0.1

PnP-Flow thus provides a favorable tradeoff between performance and resource demand, supporting large-scale or real-time applications (Martin et al., 3 Oct 2024).

4. Experimental Results and Empirical Performance

PnP-Flow is validated on challenging benchmarks (CelebA $128 \times 128$, AFHQ-Cat $256 \times 256$) across tasks:

• Gaussian denoising
• Wiener deblurring
• $2\times$, $4\times$ super-resolution
• Random (70%) and box inpainting

Performance comparison (CelebA, PSNR):

Method/Task	Denoise	Deblur	SR	Rand IP	Box IP
Degraded	20.00	27.67	7.53	11.82	22.12
PnP-Diff	31.00	32.49	31.20	31.43	—
OT-ODE	30.50	32.63	31.05	28.36	28.84
D-Flow	26.42	31.07	30.75	33.07	29.70
Flow-Priors	29.26	31.40	28.35	32.33	29.40
PnP-Flow	32.45	34.51	31.49	33.54	30.59

PnP-Flow ranks first or second in all tasks, and is uniquely able to address both classical inverse and highly generative settings with artifact-free outputs. Qualitative differences are also observed: D-Flow displays hallucinations, Flow-Priors introduces noise/textural artifacts, while PnP-Flow achieves realistic reconstructions (Martin et al., 3 Oct 2024).

With an accelerated SDE-informed variant (“IPnP-Flow”), further gains are realized, including improvements of +0.8–2.1 dB PSNR and +0.02–0.04 SSIM on standard benchmarks under identical compute budgets (Jia et al., 3 Dec 2025).

5. Advanced Variants and Applications

PnP-Flow is incorporated into several advanced schemes:

• Active learning for radio map construction, where the generative capacity is used for uncertainty quantification by producing ensembles of candidate reconstructions, guiding data acquisition for UAV navigation (Sun et al., 17 Sep 2025).
• Time-adaptive warm-up and sharp Gaussianity regularization (inspired by FMPlug) further stabilize latent optimization and keep the solution close to the high-density region of the FM prior (Wan et al., 1 Aug 2025).
• Plug-and-play priors using rectified flow models allow efficient optimization for text-to-3D generation and image inversion. Such priors support low-overhead, invertible, and time-symmetric operations, broadening the scope of PnP-based plug-in frameworks (Yang et al., 5 Jun 2024).

6. Insights, Limitations, and Open Directions

Strengths:

• Unifies and generalizes PnP and FM frameworks for a universal restoration algorithm.
• High computational/memory efficiency, especially compared with gradient-based methods requiring backpropagation through generative flows.
• Robust to initialization and able to use arbitrary latent distributions (Gaussian, Dirichlet, categorical).
• Theoretically grounded using SDEs, enabling principled improvement via schedule optimization and regularization (Jia et al., 3 Dec 2025).

Limitations:

• Reconstructions may be over-smoothed—characteristic of MMSE estimators—potentially lacking high-frequency detail.
• Most effective when the FM model induces straight-line (e.g., OT-FM or rectified) flows; performance may degrade with highly curved flows.

Open problems include:

• Extensions to non-Gaussian noise models,
• Joint fine-tuning of the FM prior within the PnP loop,
• Generalization of the reprojection mechanism to more flexible generative ODE/SDE models,
• Improved uncertainty estimation and active data acquisition strategies,
• Further acceleration and theoretical analysis via the SDE framework (Martin et al., 3 Oct 2024, Jia et al., 3 Dec 2025, Sun et al., 17 Sep 2025).

7. Summary of Key Literature

Article Title	arXiv ID	Main Contribution
PnP-Flow: Plug-and-Play Image Restoration with Flow Matching	(Martin et al., 3 Oct 2024)	Foundational algorithm; image restoration pipeline
FMPlug: Plug-In Foundation Flow-Matching Priors for Inverse Problems	(Wan et al., 1 Aug 2025)	Time-adaptive path mixing and sharp Gaussianity for domain-agnostic priors
Plug-and-Play Image Restoration with Flow Matching: A Continuous Viewpoint	(Jia et al., 3 Dec 2025)	Continuous-time SDE analysis; accelerated variant
Text-to-Image Rectified Flow as Plug-and-Play Priors	(Yang et al., 5 Jun 2024)	Use of rectified flow models as efficient, invertible priors
Flow Matching-Based Active Learning for Radio Map Construction with Low-Altitude UAVs	(Sun et al., 17 Sep 2025)	Application to uncertainty-driven active learning

PnP-Flow represents a modular, efficient, and empirically robust paradigm bridging classical optimization and deep generative modeling for inverse imaging, with continuous innovation in theoretical analysis and downstream applications (Martin et al., 3 Oct 2024, Jia et al., 3 Dec 2025, Wan et al., 1 Aug 2025, Yang et al., 5 Jun 2024, Sun et al., 17 Sep 2025).