Plug-and-Play (PnP) methods using continuous-time deep generative models as learned data distribution priors enable synthesizing high-quality solutions of model-based inverse problems. The quality of the solutions depends critically on the interplay and consensus between the inverse problem solver and the generative model sampler. We revisit a recently introduced flow-based PnP framework, which we refer to as PnP-CF (Consensus Flow), and extend its application to nonlinear inverse problems, including those with learned neural-network forward operators. PnP-CF combines a flow-matching denoiser with (i) ADMM-like dual variable updates and (ii) a time-adaptive pullback onto the intermediate flow probability path that combines a stochastic and a deterministic component. We discuss how these two features stabilize the algorithm and promote consensus, and demonstrate on a set of experiments that PnP-CF not only solves standard linear inverse problems with quality at least comparable to state-of-the-art baselines, but also extends gracefully to nonlinear inverse problems, including forward operators realized by pretrained neural networks, for which existing methods produce unsatisfactory results or fail completely.
Figures & tables
Figure 1: Phases of PnP-CF. Left: Starting from z~ , the denoiser D(z~) produces z^ (green). The dual correction (gray) shifts z^ by u , and the forward-consistency operator F(y,x~) maps it to x^ (blue), pulling it towards the region consistent with y under l(y;x) . Right: The time-adaptive pullback constructs the next z~ . The deterministic component ϵd , computed from the previous z~ and z^ , combines with ϵs∼pt=0 to form ϵ . The dual-corrected point x^+u (gray) and the ϵ together produces the new z~ , which serves as the input to the denoiser at the next iteration.
Figure 2: Ablation of PnP-CF update strategies on the box-inpainting inverse problem. Left: qualitative comparison of stochastic-only, deterministic-only, and time-adaptive pullback strategies. Right: effect of the dual variable u on the forward-consistency and denoising updates, shown through PSNR and SSIM trajectories over iterations, averaged over 100 images.
Figure 3: Evolution of the primal residual ∥x^k−z^k∥2 across linear and nonlinear inverse problems on AFHQ-Cat (left) and CelebA (right).
Figure 4: Comparison of image restoration methods on box inpainting, super-resolution, and random inpainting, with additional qualitative results provided in appendix F . Our method yields visually sharp reconstructions while preserving structural consistency across tasks.
Method
Time (s)
Mem (GB)
PnP-CF (ours)
17.62
0.30
PnP-Flow
16.40
0.30
D-Flow
71.42
13.34
FlowDPS
11.14
0.32
Table 2: Per batch runtime and peak GPU memory (HDR recon. - CelebA)
Figure 5: Qualitative results. Left: Segmentation-guided reconstruction results for D-Flow and PnP-CF. Right: Qualitative reconstruction on nonlinear inverse problems, including HDR reconstruction, phase retrieval, and nonlinear deblurring with additional results in appendix G .
Table 4: Hyperparameter settings per task and dataset for the time-adaptive PnP-CF variant. K : number of PnP-CF iterations ( steps_pnp ); ρ : base ADMM penalty ( rho_admm ); β : penalty-schedule exponent ( beta_rho ); γ : time-schedule exponent ( t_schedule_gamma ); Denoiser calls: number of denoiser samples averaged per iteration ( num_denoiser_calls ); GD steps: nonlinear x-update gradient steps ( admm_gd_steps , linear tasks use CG instead).
Box Inpaint.
Super-res.
Rand. Inpaint.
HDR Recon.
Phase Retr.
Nonl. Deblur.
OT-ODE
t0 (start time)
0.1
0.1
0.1
N/A
N/A
N/A
γ schedule
γt
constant
constant
N/A
N/A
N/A
K (steps)
100
100
100
N/A
N/A
N/A
Flow Priors
λ (regularization)
10,000
10,000
10,000
N/A
N/A
N/A
Appendix
Table 5: Hyperparameters used for all baseline methods on the CelebA dataset.
Box Inpaint.
Super-res.
Rand. Inpaint.
HDR Recon.
Phase Retr.
Nonl. Deblur.
OT-ODE
t0 (start time)
0.1
0.1
0.1
N/A
N/A
N/A
γ schedule
γt
constant
constant
N/A
N/A
N/A
K (steps)
100
500
100
N/A
N/A
N/A
Flow Priors
λ (regularization)
10,000
10,000
10,000
N/A
N/A
N/A
Appendix
Table 6: Hyperparameters used for all baseline methods on the AFHQ-Cat dataset.
Figure 6: Effect of the dual-variable update on reconstruction quality across optimization iterations. PSNR (top), SSIM (middle), and LPIPS (bottom) are shown for super-resolution, random inpainting, phase retrieval, and HDR reconstruction on CelebA, comparing PnP-CF ‘with’ and ‘without’ the dual-variable update. Across tasks, the dual-variable formulation generally improves convergence speed and final reconstruction quality, with particularly large gains for super-resolution, random inpainting, and HDR reconstruction.
CelebA
AFHQ-Cat
Task
Method
PSNR ↑
LPIPS ↓
SSIM ↑
PSNR ↑
LPIPS ↓
SSIM ↑
HDR Reconstruction
PnP-CF (adaptive)
27.59
0.04
0.85
26.29
0.08
0.86
PnP-CF (stochastic)
22.44
0.10
0.80
22.97
0.32
0.71
PnP-CF (deterministic)
18.12
0.36
0.43
24.11
0.13
0.75
Box Inpainting
PnP-CF (adaptive)
26.36
0.05
0.89
22.65
0.13
0.82
PnP-CF (stochastic)
22.79
0.18
0.74
21.24
0.37
0.68
Appendix
Table 8: Task-wise performance of the three PnP-CF pullback variants (time-adaptive, stochastic, deterministic) across all six inverse tasks. Best result per task and metric, among the three variants, is shown in bold .
Figure 7: Additional qualitative ablation results comparing the pullback variants of PnP-CF on CelebA (left) and AFHQ-Cat (right). Within each dataset, columns from left to right show the clean image, degraded observation, stochastic-only ( ϵs ), deterministic-only ( ϵd ), and time-adaptive ( ϵs,ϵd ) variants. From top to bottom, the rows correspond to box inpainting, random inpainting, super-resolution, phase retrieval, nonlinear deblurring, and HDR reconstruction.
Figure 8: Failure cases for segmentation-guided reconstruction. Left: reference images. Middle: PnP-CF reconstructions. Right: reconstructed images with the predicted segmentation boundaries overlaid in red. Despite achieving strong segmentation consistency, some reconstructions exhibit distorted facial geometry and perceptual artifacts. These examples illustrate that satisfying a learned segmentation forward model does not necessarily guarantee a visually plausible reconstruction.
PSNR ↑
SSIM ↑
LPIPS ↓
Dataset
Task
Rec
Noisy
Rec
Noisy
Rec
Noisy
CelebA
Box Inpainting
24.71
18.91
0.76
0.66
0.08
0.31
Random Inpainting
30.15
11.46
0.90
0.18
0.04
1.09
Super Resolution
20.28
10.71
0.56
0.18
0.24
0.82
HDR Reconstruction
21.97
17.03
0.79
0.50
0.06
0.11
Nonlinear Deblurring
25.36
24.77
0.73
0.66
0.15
0.24
Appendix
Table 9: PnP-CF results on reconstructed vs. noisy/degraded inputs with default (non-scheduled) hyperparameter settings on CelebA and AFHQ-Cat.
Box inpainting
Super-resolution
Random inpainting
Dataset
Method
PSNR ↑
LPIPS ↓
SSIM ↑
PSNR ↑
LPIPS ↓
SSIM ↑
PSNR ↑
LPIPS ↓
SSIM ↑
CelebA
PnP-CF (ours)
31.71
0.02
0.95
31.37
0.03
0.86
33.99
0.02
0.95
PnP-Flow
30.82
0.04
0.94
31.31
0.06
0.91
33.92
0.02
0.95
FLOWER
31.66
0.02
0.95
32.88
0.04
0.90
33.91
0.03
0.95
FlowDPS
28.48
0.06
0.89
25.94
0.12
0.79
26.42
0.11
0.80
D-Flow
28.28
0.06
0.85
29.40
0.04
0.83
31.96
0.03
0.91
Appendix
Table 10: Reconstruction quality on CelebA and AFHQ-Cat for linear inverse problems under the milder degradation settings used in Martin et al. [2025] . Best results are shown in bold .
Figure 9: Additional qualitative comparisons on the CelebA dataset across linear inverse tasks, including box inpainting (1st row), random inpainting (2nd row) and super-resolution (3rd row).
Figure 10: Additional qualitative comparisons on the AFHQ-Cat dataset across linear inverse tasks, including box inpainting (1st row), random inpainting (2nd row) and super-resolution (3rd row).
Figure 11: Additional qualitative comparisons on the CelebA dataset across nonlinear inverse tasks, including phase retrieval (1st row), nonlinear deblurring (2nd row) and HDR reconstruction (3rd row).
Figure 12: Additional qualitative comparisons on the AFHQ-Cat dataset across nonlinear inverse tasks, including phase retrieval (1st row), nonlinear deblurring (2nd row) and HDR reconstruction (3rd row).
Figure 13: Qualitative comparison for segmentation-guided reconstruction on CelebA-HQ. From left to right: reference image, target segmentation mask, and reconstructions from PnP-Flow, FlowDPS, D-Flow, and PnP-CF. PnP-Flow exhibits mode collapse across different masks, while FlowDPS produces plausible but weakly conditioned reconstructions. D-Flow follows the segmentation constraint more closely but introduces severe visual artifacts. In contrast, PnP-CF produces visually plausible reconstructions while maintaining strong consistency with the target segmentation.
Plug-and-play (PnP) methods for solving inverse problems have recently achieved strong performance by leveraging denoising priors based on powerful generative diffusion and flow models. However, existing diffusion- and flow-based PnP methods typically rely on stochastic renoise-denoise operations, which complicate the analysis of their convergence behavior. In this work, we identify and formalize the deterministic renoise-denoise operator underlying flow-based plug-and-play methods. This perspective reveals that these methods implicitly define a deterministic operator given by the expectation of a denoiser over the latent noise distribution. Building on this insight, we propose FlowADMM, a PnP algorithm that integrates the renoise-denoise operator into the classical alternating direction method of multiplier (ADMM) framework. We establish convergence guarantees for FlowADMM under weak Lipschitz conditions on the underlying flow network, and extend the analysis to non-stationary time schedules. Empirically, FlowADMM achieves state-of-the-art performance among flow-based PnP methods on a range of inverse problems, including denoising, deblurring, super-resolution, and inpainting, while requiring fewer data consistency evaluations than prior approaches.
Hendrik Sommerhoff, Michael Moeller
Computer Vision Group, University of Siegen Hölderlinstraße 3, 57076 Siegen, Germany
Plug-and-play (PnP) methods are widely used for solving imaging inverse problems by incorporating a denoiser into optimization algorithms. Score-based diffusion models (SBDMs) have recently demonstrated strong generative performance through a denoiser trained across a wide range of noise levels. Despite their shared reliance on denoisers, it remains unclear how to systematically use SBDMs as priors within the PnP framework without relying on reverse diffusion sampling. In this paper, we establish a score-based interpretation of PnP that justifies using pretrained SBDMs directly within PnP algorithms. Building on this connection, we introduce a stochastic generative PnP (SGPnP) framework that injects noise to better leverage the expressive generative SBDM priors, thereby improving robustness in severely ill-posed inverse problems. We provide a new theory showing that this noise injection induces optimization on a Gaussian-smoothed objective and promotes escape from strict saddle points. Experiments on challenging inverse tasks, such as multi-coil MRI reconstruction and large-mask natural image inpainting, demonstrate consistent improvement over conventional PnP methods and achieve performance competitive with diffusion-based solvers.
Chicago Y. Park, Edward P. Chandler, Yuyang Hu +4
University of Wisconsin–Madison · Washington University in St. Louis · Los Alamos National Laboratory
Flow matching has emerged as the state-of-the-art generative model and has been used for plug-and-play (PnP) priors to solve inverse problems in computational imaging. However, existing flow-based inverse solvers make simplifying approximations in posterior sampling and/or assume linear forward models. In contrast, principled diffusion-based methods exist but are computationally expensive. To circumvent these problems, we introduce FlowSGS, a flow-based posterior sampling method using Split Gibbs Sampling (SGS) with annealed coupling, and make flow-specific interventions to the prior sampling and annealing steps, theoretically grounded in the Stochastic Interpolants (SI) framework. Using a novel timestep correction and accelerated warm-up with one-step flow models, our method is faster and more accurate than the state of the art across a range of imaging problems. We demonstrate the first use of flow-based samplers for a nonlinear inverse problem (Fourier phase retrieval) and apply one-step flow models to principled Bayesian inverse imaging.
Tianao Li, Xinhui Qian, Emma Alexander
Department of Computer Science, Northwestern University · NSF-Simons AI Institute for the Sky (SkAI) · Department of Statistics and Data Science, Northwestern University