Technical Summary: FlowADMM: Plug-and-play ADMM with Flow-based Renoise-Denoise Priors

Problem Statement

Inverse problems in imaging, such as denoising, deblurring, super-resolution, and inpainting, are typically ill-posed and formulated as recovering an image $x$ from degraded measurements $y = Ax + \nu$. While deep learning offers powerful direct mappings, they often lack interpretability and theoretical guarantees. Classical variational approaches formulate this as minimizing a data fidelity term plus a regularization term, often solved via proximal algorithms like Forward-Backward Splitting (FBS).

Plug-and-Play (PnP) methods improve upon this by replacing the proximal operator of a hand-crafted regularizer with a learned denoiser. Recently, generative flow and diffusion models have provided time-dependent families of denoisers $\{D_t\}$ adapted to specific noise levels. However, a key challenge arises: during iterative reconstruction, the intermediate iterates often deviate from the noise distribution $p_t$ for which the denoiser was trained. Existing flow-based PnP methods (e.g., PnP-Flow, Flower) address this by incorporating explicit renoise-denoise steps, where the current iterate is perturbed toward the appropriate noise level before denoising. These methods often average over multiple noise samples to stabilize the process.

Despite their empirical success, existing renoise-denoise PnP methods are formulated as stochastic algorithms, making their convergence behavior difficult to analyze theoretically. Specifically, the deterministic operator implicitly induced by averaging the denoiser output over the latent noise distribution has not been characterized, leaving the convergence properties of these methods poorly understood.

Methodology

The Mean Renoise-Denoise Operator

The authors identify and formalize the deterministic operator underlying stochastic flow-based PnP methods. Given a pretrained flow velocity network $v_\theta$ defining a time-dependent MMSE denoiser $D_t = I + (1-t)v_\theta^t$, the standard stochastic step involves renoising an iterate $x$ with noise $\epsilon \sim \mathcal{N}(0, I)$ to get $tx + (1-t)\epsilon$, followed by denoising.

The paper defines the mean renoise-denoise operator $\bar{S}_t(x)$ as the expectation of this process over the noise distribution:
$$\bar{S}_t(x) = \mathbb{E}_{\epsilon \sim \mathcal{N}(0, I)} [D_t(tx + (1-t)\epsilon)]$$
This operator represents the deterministic limit of the stochastic updates used by methods like PnP-Flow. The authors analyze its residual form $R_t(x) = \bar{S}_t(x) - x$, interpreting it as a vector field providing local corrections toward the data manifold. Under ideal MMSE assumptions, points on the data manifold are fixed points of this operator.

FlowADMM Algorithm

Building on this insight, the authors propose FlowADMM, which integrates the mean renoise-denoise operator into the Alternating Direction Method of Multipliers (ADMM) framework. Unlike FBS, ADMM allows for multiple prior operators and handles non-differentiable data terms without operator-dependent step-size restrictions.

The algorithm solves the constrained problem $\min_{x,z} F_y(x) + R(z)$ subject to $x=z$ by minimizing the augmented Lagrangian. The updates are:

1. Data Consistency: $x^{k+1} = \text{prox}_{\tau F_y}(z^k - u^k)$
2. Prior Update: $z^{k+1} = \bar{S}_{t_k}(x^{k+1} + u^k)$
3. Dual Update: $u^{k+1} = u^k + x^{k+1} - z^{k+1}$

Since the expectation in $\bar{S}_t$ is intractable, it is approximated via Monte Carlo sampling. The authors propose a timestep-dependent sampling strategy: using fewer samples in early iterations (where coarse corrections are made and the operator is less sensitive to noise) and increasing the number of samples in later iterations (where fine details are refined and precision is critical).

Convergence Analysis

The paper establishes convergence guarantees for FlowADMM under weak Lipschitz conditions on the underlying flow network:

• Fixed Time Schedule: For a constant timestep $t$, convergence is guaranteed if the flow network's Lipschitz constant $L_v$ satisfies $L_v < \frac{1}{1-t}$. This condition is significantly weaker than the strict contraction ($L < 1$) required by many other PnP methods, especially as $t \to 1$.
• Non-Stationary Time Schedules: The analysis is extended to time schedules $(t_k)_k$ that converge to a limit $t_{max} < 1$. The authors prove convergence to a fixed point of the limiting operator $T_{t_{max}}$ provided the sequence of operators approaches the limit sufficiently fast and the flow network is Lipschitz continuous with respect to time $t$.

Key Contributions

1. Theoretical Characterization: The authors identify and analyze the deterministic mean renoise-denoise operator underlying stochastic flow-based PnP methods, characterizing its residual structure and Lipschitz behavior.
2. Algorithm Proposal: They propose FlowADMM, a PnP algorithm that integrates this mean operator directly into the ADMM framework.
3. Convergence Guarantees: They derive convergence criteria for FlowADMM under weak Lipschitz assumptions, covering both constant and non-constant (time-varying) timestep schedules.
4. Sampling Strategy: They propose a coarse-to-fine Monte Carlo sampling strategy that increases the number of samples in later iterations, motivated by the behavior of the mean operator.
5. Empirical Performance: They demonstrate state-of-the-art performance across multiple inverse problems while requiring fewer data consistency evaluations than prior stochastic averaging approaches.

Experimental Results

The authors evaluated FlowADMM on the CelebA and AFHQ-Cat datasets across five tasks: denoising, deblurring, super-resolution, random inpainting, and box inpainting.

• Performance: FlowADMM achieved state-of-the-art results in distortion metrics (PSNR and SSIM) on 7 out of 10 tasks compared to recent flow and diffusion-based PnP methods (including PnP-Flow, Flower, OT-ODE, D-Flow, and DiffPIR). Notable improvements were observed in super-resolution and inpainting tasks.
• Efficiency: While matching the total number of flow evaluations (network forward passes) of the leading competitor Flower5-OT, FlowADMM required significantly fewer data-term proximal operator evaluations. This is because Flower5-OT averages full reconstruction trajectories, whereas FlowADMM averages only the prior step within each iteration.
• Ablation Studies: Experiments confirmed that the "late-heavy" sampling schedule (increasing samples in later iterations) significantly improves performance over constant sampling. Furthermore, applying this schedule to PnP-Flow improved its results, but FlowADMM still outperformed it, suggesting the ADMM formulation itself provides an advantage.
• Limitations: The convergence analysis relies on Lipschitz constants of the pretrained network which are not explicitly penalized during training. Additionally, the theoretical guarantees apply to the deterministic mean operator, while the practical implementation uses stochastic Monte Carlo approximations.

Significance and Claims

The paper claims that by explicitly formalizing the deterministic operator underlying stochastic renoise-denoise methods, it bridges the gap between empirical success and theoretical understanding in flow-based PnP. The work provides:

• Provable Convergence: Unlike previous stochastic PnP methods, FlowADMM comes with convergence guarantees under realistic Lipschitz conditions on the flow network.
• Algorithmic Efficiency: The integration into ADMM allows for a more efficient use of computational resources (fewer data consistency steps) compared to trajectory-averaging methods.
• Operator-Theoretic Insight: The analysis suggests that the "coarse-to-fine" behavior of the mean operator can be exploited through adaptive sampling schedules, a principle that benefits both ADMM and FBS-based approaches.

The authors conclude that FlowADMM offers a robust, theoretically grounded framework for solving inverse problems with flow-based priors, though they acknowledge that future work could focus on enforcing Lipschitz constraints during training to strengthen stability guarantees.