P-Flow: Proxy-gradient Flows for Linear Inverse Problems
For practitioners of inverse problems (e.g., image restoration), P-Flow offers a more stable and efficient flow-based generative approach, though it is an incremental improvement over existing flow matching methods.
P-Flow introduces a proxy gradient method to stabilize and accelerate flow matching for linear inverse problems, avoiding numerical instability and memory overhead of unrolled differentiation. It achieves competitive performance, particularly under extreme degradations like severe ill-posedness and high noise.
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.