Flow Priors for Linear Inverse Problems via Iterative Corrupted Trajectory Matching

2024-05-29

arXiv.org

AI Summary

Flow matching (FM) generative models face high computational costs when estimating image log-likelihoods, hindering their efficient application to large-scale linear inverse problems—including super-resolution, deblurring, inpainting, and compressed sensing. Method: This paper proposes an efficient maximum a posteriori (MAP) estimation framework leveraging FM priors. It introduces two key innovations: (1) an iterative corruption trajectory matching algorithm that decomposes the global MAP optimization into sequentially solvable local subproblems, with rigorous theoretical equivalence proof; and (2) the first integration of the Tweedie formula into the FM framework, enabling gradient updates without solving ordinary differential equations (ODEs) or backpropagating through ODE solvers. Results: Extensive experiments demonstrate that our method significantly outperforms existing FM-based baselines across multiple inverse problems, achieving superior trade-offs between reconstruction quality and inference speed. The implementation is publicly available.

Technology Category

• Search and Optimization: Learning to Search
• Machine Learning: Matrix & Tensor Methods
• Computer Vision: Learning & Optimization for CV

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

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.

Problem

• Computational_Flow_Matching
• Image_Processing
• Efficiency_Optimization

Innovation

• Cyclic Algorithm
• Image Processing
• Efficiency Improvement