P-Flow: Proxy-gradient Flows for Linear Inverse Problems

2026-05-08

AI Summary

This work addresses the numerical instability and high computational cost arising from long-chain differentiation in flow-matching-based generative models when solving linear inverse problems. To overcome these challenges, the authors propose the P-Flow framework, which introduces a proxy gradient to update the source point without directly differentiating through the trajectory. Additionally, Gaussian spherical projection is incorporated to preserve consistency with the prior distribution, thereby enhancing both reconstruction stability and computational efficiency. The theoretical analysis integrates Bayesian inference, Lipschitz continuity, and concentration of measure phenomena in high dimensions. Experimental results demonstrate that P-Flow achieves superior performance across various image restoration tasks, exhibiting remarkable robustness and competitiveness even under severely ill-posed conditions and high noise levels.

Abstract

Generative models based on flow matching have emerged as a powerful paradigm for inverse problems, offering straighter trajectories and faster sampling compared to diffusion models. However, existing approaches often necessitate differentiating through unrolled paths, leading to numerical instability and prohibitive computational overhead. To address this, we propose P-Flow, a framework that stabilizes the reconstruction process by leveraging a proxy gradient to update the source point. This approach effectively circumvents the numerical instability and memory overhead of long-chain differentiation. To ensure consistency with the prior distribution, we employ a Gaussian spherical projection motivated by the concentration of measure phenomenon in high-dimensional spaces. We further provide a theoretical analysis for P-Flow based on Bayesian theory and Lipschitz continuity. Experiments across diverse restoration tasks demonstrate that P-Flow delivers competitive performance, especially under extreme degradations such as severely ill-posed conditions and high measurement noise.

Problem

• linear inverse problems
• flow matching
• numerical instability
• computational overhead
• generative models

Innovation

• proxy gradient
• flow matching
• linear inverse problems
• Gaussian spherical projection
• numerical stability