Generative models based on flow matching have attracted significant attention for their simplicity and superior performance in high-resolution image synthesis. By leveraging the instantaneous change-of-variables formula, one can directly compute image likelihoods from a learned flow, making them enticing candidates as priors for downstream tasks such as inverse problems. In particular, a natural approach would be to incorporate such image probabilities in a maximum-a-posteriori (MAP) estimation problem. A major obstacle, however, lies in the slow computation of the log-likelihood, as it requires backpropagating through an ODE solver, which can be prohibitively slow for high-dimensional problems. In this work, we propose an iterative algorithm to approximate the MAP estimator efficiently to solve a variety of linear inverse problems. Our algorithm is mathematically justified by the observation that the MAP objective can be approximated by a sum of $N$ ``local MAP'' objectives, where $N$ is the number of function evaluations. By leveraging Tweedie's formula, we show that we can perform gradient steps to sequentially optimize these objectives. We validate our approach for various linear inverse problems, such as super-resolution, deblurring, inpainting, and compressed sensing, and demonstrate that we can outperform other methods based on flow matching. Code is available at https://github.com/YasminZhang/ICTM.
翻译:基于流匹配的生成模型因其在高分辨率图像合成中的简洁性和卓越性能而备受关注。通过利用瞬时变量变换公式,可以直接从学习到的流中计算图像似然,使其成为下游任务(如逆问题)中极具吸引力的先验候选。特别地,一种自然的方法是将此类图像概率纳入最大后验(MAP)估计问题中。然而,主要障碍在于对数似然计算的缓慢性,因为这需要通过常微分方程(ODE)求解器进行反向传播,对于高维问题而言可能极其耗时。在本工作中,我们提出了一种迭代算法来高效逼近MAP估计量,以解决各类线性逆问题。我们算法的数学依据基于以下观察:MAP目标函数可通过$N$个“局部MAP”目标函数的和来近似,其中$N$为函数评估次数。通过利用Tweedie公式,我们证明了可以通过梯度步骤顺序优化这些目标。我们在多种线性逆问题(如超分辨率、去模糊、修复和压缩感知)上验证了所提方法,并证明其性能优于其他基于流匹配的方法。代码发布于https://github.com/YasminZhang/ICTM。