Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Flow Priors for Linear Inverse Problems via Iterative Corrupted Trajectory Matching

Yasi Zhang, Peiyu Yu, Yaxuan Zhu, Yingshan Chang, Feng Gao, Ying Nian Wu, Oscar Leong

TL;DR

This work tackles linear inverse problems by leveraging flow-based priors within a MAP framework. It introduces Iterative Corrupted Trajectory Matching (ICTM), an algorithm that decomposes the global MAP objective into a sequence of local objectives over an interpolated trajectory, enabling gradient-based optimization via Tweedie's formula without backpropagation through full ODE solvers. The method is validated on natural (CelebA-HQ) and medical (HCP T2w) datasets, showing consistent improvements over flow-based and diffusion-based baselines for tasks including super-resolution, deblurring, inpainting, and compressed sensing. While demonstrating substantial practical gains, the work discusses limitations related to trajectory compliance, extension to nonlinear forward models, and the need for uncertainty quantification in reconstructions.

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.

Flow Priors for Linear Inverse Problems via Iterative Corrupted Trajectory Matching

TL;DR

This work tackles linear inverse problems by leveraging flow-based priors within a MAP framework. It introduces Iterative Corrupted Trajectory Matching (ICTM), an algorithm that decomposes the global MAP objective into a sequence of local objectives over an interpolated trajectory, enabling gradient-based optimization via Tweedie's formula without backpropagation through full ODE solvers. The method is validated on natural (CelebA-HQ) and medical (HCP T2w) datasets, showing consistent improvements over flow-based and diffusion-based baselines for tasks including super-resolution, deblurring, inpainting, and compressed sensing. While demonstrating substantial practical gains, the work discusses limitations related to trajectory compliance, extension to nonlinear forward models, and the need for uncertainty quantification in reconstructions.

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 ``local MAP'' objectives, where 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.
Paper Structure (43 sections, 6 theorems, 41 equations, 14 figures, 3 tables, 1 algorithm)

This paper contains 43 sections, 6 theorems, 41 equations, 14 figures, 3 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background
  3. Notation
  4. Flow-Based Models
  5. Probability Computation for Flow Priors
  6. Method
  7. Flow-Based MAP
  8. Flow-Based MAP Approximation
  9. Discussion of assumptions:
  10. Local data likelihood:
  11. Local prior:
  12. Toy Example Validation
  13. Experiments
  14. Baselines
  15. Natural Images
  16. ...and 28 more sections

Key Result

Theorem 1

For $N \geq 1$, set $\gamma_i := (\frac{1}{2})^{N-i+1}$ and $\Delta t = 1/N$. Suppose $y = \mathcal{A}(x_*) + \epsilon$ where $x_* = x_1(x_0)$ with $x_0$ being the solution to Eq. eq:globalupdate, $\epsilon \sim \mathcal{N}(0,\sigma_y^2I)$, and $x_t$ exactly follows the straight path $x_t = tx + (1 where $\hat{\mathcal{J}}_i = \log p(x_{(i-1) \Delta t}) -\mathop{\mathrm{tr}}\left (\frac{\partial

Figures (14)

  • Figure 1: Illustration of the idea of ICTM. The corrupted trajectory $u_t: = \mathcal{A}(x_t)$ follows the corrupted flow ODE $du_t = \mathcal{A}(v_{\theta}(x_t,t))dt.$
  • Figure 2: Results of a toy example modeling 1,000 FFHQ faces as a Gaussian distribution. Subfigure (a) shows the qualitative results of our method; Subfigure (b) presents the histogram of the differences between ours and the true MAP; Subfigure (c) displays the MSE values as the NFEs varies.
  • Figure 3: Qualitative comparison results on the CelebA-HQ dataset. The reconstructions generated by our method align more faithfully with the ground truth and exhibit a higher degree of refinement.
  • Figure 4: Qualitative comparison results on compressed sensing. Our method produces more faithful reconstructions with fewer artifacts, ensuring higher accuracy and clarity in the details.
  • Figure 5: Ablation results of step size $\eta$ and guidance weight $\lambda$. The choice of hyperparameters for our algorithm is fairly consistent across all tasks. We choose $\eta = 10^{-2}$ for all experiments on CelebA-HQ. For $\lambda$, we choose $\lambda = 10^3$ for Gaussian deblurring and $\lambda = 10^4$ for the other tasks.
  • ...and 9 more figures

Theorems & Definitions (9)

  • Theorem 1
  • Proposition 1
  • Lemma 1
  • proof
  • Lemma 2: Tweedie's Formula efron2011tweedie
  • Lemma 3
  • proof
  • Proposition 2
  • proof