FlowSGS improves image recovery from complex measurements using flow models

FlowSGS: Improving Flow Matching Priors for Inverse Imaging with Stochastic Interpolants

Computer Vision and Pattern Recognition

Summary

Recovering original images from incomplete or distorted data is a common problem in imaging. The authors present FlowSGS, a new method that better combines knowledge about the image and the measurement process to reconstruct clear images. Unlike earlier methods, it can handle complex, nonlinear problems and requires fewer steps to produce high-quality results. This makes FlowSGS a promising tool for solving tricky image restoration tasks.

What this means in practice

Authors

Tianao Li, Xinhui Qian, Emma Alexander

Abstract

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 assume linear forward models and/or make simplifying approximations in posterior sampling. To circumvent these problems, we introduce FlowSGS, a flow-based posterior sampling method using Split Gibbs Sampling (SGS) to decompose the posterior into a likelihood step and a prior step. Specifically, we sample from the likelihood step using Langevin dynamics and leverage the Stochastic Interpolants (SI) framework to integrate a pretrained flow model into the prior step. We provide a form for the prior step that uses SI's reverse-time SDE, and show connections to previous PnP methods. Moreover, with the aid of the flow prior's straight probability paths and a novel timestep correction technique for the reverse-time SDE, FlowSGS requires fewer network evaluations in its prior step than plug-and-play diffusion samplers. Our experiments show state-of-the-art performance on a range of inverse problems. For the first time, we provide an experiment on a nonlinear inverse problem (Fourier phase retrieval) for flow-based inverse solvers.