Flow Priors for Linear Inverse Problems via Iterative Corrupted Trajectory Matching

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.

Paper

References (68)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer qjx25/10 · confidence 4/52024-07-11

Summary

The paper investigates the application of flow matching-based generative models for high-resolution image synthesis, particularly as priors for solving inverse problems. A notable challenge addressed is the slow computation of log-likelihoods in high-dimensional contexts, necessitating backpropagation through an ODE solver. To overcome this, the authors propose an efficient iterative algorithm for approximating the maximum-a-posteriori (MAP) estimator. This method involves approximating the MAP objective through a series of "local MAP" objectives and employs Tweedie's formula for sequential gradient optimization. The proposed method's performance is validated across multiple inverse problems and various datasets.

Strengths

- The paper introduces a unique method for solving inverse problems and reconstructing a single image. - The proposed method is supported by a solid theoretical foundation. - The approach is effective across diverse linear inverse problems.

Weaknesses

- Insufficient Empirical Evidence: The paper lacks enough qualitative results and empirical evidence to conclusively demonstrate the superiority of the proposed method. - Need for Additional Metrics: Metrics like FID and LPIPS scores should be included alongside PSNR and SSIM. - Figure 2 Presentation: The data in Figure 2 should be presented in a table format for easier interpretation of quantitative performance improvements. - Limited Dataset Testing: For natural image experiments, the algorithm was only tested on the CelebA-HQ dataset. Pretrained models for other datasets (LSUN-bedroom, LSUN-church, and AFHQ-cat) are publicly available from the authors of the rectified flow paper and should be included. - Blurry reconstruction results: A closer look at Figure 3 (a,c) reveals that the reconstructed images are too smooth and lack high-frequency details. Compared to OT-ODE, the results are smoother and blurrier. - Comparison with Recent Work: Apart from OT-ODE, recent work like "D-Flow: Differentiating through Flows for Controlled Generation" also addresses inverse problems and should be considered for comparison.

Questions

- Does the proposed method work for very noisy corruptions with high measurement noise (e.g., sigma_y = 0.2 or more)? - Does the proposed method work for severely ill-posed inverse problems like 8x super-resolution or inpainting with 90% of pixels missing? - Why was the algorithm evaluated on only 100 images from the CelebA-HQ dataset? Why wasn't the entire validation dataset used? -There are numerous qualitative results for medical images in the appendix, but for natural image datasets like CelebA-HQ, there is only one small figure (Figure 3). It is hard to draw any conclusions based on just Figure 3. Was the authors' focus primarily on medical images?

Rating

5

Confidence

4

Soundness

3

Presentation

2

Contribution

2

Limitations

- The paper's empirical validation is limited, requiring more comprehensive testing across various datasets and additional qualitative results to strengthen the claims of superiority. - The presentation of quantitative results could be improved by using tables for easier comparison and interpretation. - The evaluation metrics need to be expanded to include FID and LPIPS scores to provide a more comprehensive assessment of performance. - Blurry results

Authorsrebuttal2024-08-12

Dear reviewer, Thanks a lot for your time and effort. As the discussion period approaches its deadline, we would greatly appreciate your feedback on whether our rebuttal has adequately addressed your concerns. Please reach out if you have additional questions or comments which we are happy to address. Thank you again. Best, Authors

Reviewer qjx22024-08-12

The authors have addressed my concerns by providing empirical evidence for the requested cases. Hence, I raise my score to 5.

Reviewer sT2G5/10 · confidence 4/52024-07-12

Summary

The paper proposes a flow prior under the MAP structure for solving inverse problems, a theoretical analysis is also given.

Strengths

1. The paper is easy to follow, the presentation from concepts to methods is concise, the motivation of the proposed method on overcoming shortage of flow model is also clear and straightforward. 2. Incorporating the flow prior into MAP to solve inverse problems is interesting and worth exploring. 3. The discussion on feasibility of assumptions made in the paper, as well as hyperparameter effects are very informative to the readers, and are valuable to the research in the community.

Weaknesses

1. Theorem 1 is trivial in terms of $N \rightarrow 0$, and the existence of constant $c(N)$. The approximation gap converges to 0 as N goes to infinity does not provide insightful guidance on the practical implementation since $N \rightarrow 0$ is infeasible (decreases efficiency) in practice. Also, the implicit representation of approximation error constant $c(N)$ does not make sense because it gives no information about how the error could be scaled with $N$, the authors are expected to provide explicit expression of $c( \cdot )$ as a function of $N$. 2. Obtaining measurements $u_t$ in the ‘corrupted trajectory’ and generating auxiliary path $s_t$ do not make too much sense to me, especially when $x_t$ is very noisy and the forward operator $A(\cdot)$ is challenging. Suppose $A$ is very ill-posed (e.g. compressed sensing with low sampling rate), then the $y_t$ in the auxiliary path is not good itself, nevertheless the measurement $u_t = A(x_t)$ obtained in the corrupted trajectory especially when $x_t$ has high noise level. The assumption on exact compliance of trajectory (between corrupted trajectory and auxiliary path) is not theoretically guaranteed as stated by the authors (line 170-171), and the strong empirical results can not support the existence of the assumption made. The good empirical results could only be valid in less challenging inverse problems, or those inverse problems that are suitable for generative models to solve, e.g. inpainting, super-resolution. 3. In the compressed sensing experiment, the sampling rate $\nu$ is relatively large (0.25, 0.5). The authors are expected to use smaller sampling rates (e.g. 0.05, 0.1), and at the same time compare it with classical recovery algorithms, developed from the seminal work by Donoho, et al [1,2] (no neural networks involved), as well as compressed sensing with other generative models (VAE, GAN, Diffusion). The current results and comparison are not convincing. [1] Donoho, David L. "Compressed sensing." IEEE Transactions on information theory 52.4 (2006): 1289-1306. [2] Lustig, Michael, et al. "Compressed sensing MRI." IEEE signal processing magazine 25.2 (2008): 72-82. 4. The $\lambda_t$ used in Proposition 1 as SNR can also be found in previous work [1], the authors are expected to provide some comparison or discussion with [1] although it is based on Diffusion, it uses MAP structure and solve optimization problem during the sampling process. [1] Mardani, Morteza, et al. "A variational perspective on solving inverse problems with diffusion models." arXiv preprint arXiv:2305.04391 (2023). 5. The number of iterations $K$ affects the performance a lot on different tasks, the authors are expected to provide more explanation on the choice of $K$. Also, how are the $K$ and $\lambda$ picked in the experiments? Given that the testing performances are sensitive to these hyperparameters. 6. The baseline methods compared are limited, only 3 flow based models are picked, generative models have been widely used for image restoration tasks, the authors are expected to provide at least several other representative generative models for comparison.

Questions

N/A

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

N/A

Authorsrebuttal2024-08-12

We greatly appreciate your confidence in our work and your thoughtful suggestions regarding our theoretical contributions. We promise to incorporate the necessary adjustments in the revision. **Before the discussion period ends, we would like to provide a last-minute highlight of the strengths of our approach, particularly in comparison to the concurrent flow-based method, D-Flow [1], as recognized by other reviewers.** As shown in the Table 3 in the appendix of our paper, our method requires 1.6 min for each image, OT-ODE requires 2.46 min, while the concurrent work D-Flow [1], which formulates the MAP into a constrained optimization problem (Eq. 9), requires 5-10 mins for each image as documented in their Sec. 3.4. This is because each of its optimization steps requires a backpropogation through an ODE solver, as they calculate the log-likelihood fully. Our work is significantly faster due to our principled local MAP approximation. We hope this runtime comparison further highlights the novelty of our method. [1] D-Flow: Differentiating through Flows for Controlled Generation. Heli Ben-Hamu, Omri Puny, Itai Gat, Brian Karrer, Uriel Singer, Yaron Lipman. ICML 2024 (concurrent work)

Reviewer Xm5z7/10 · confidence 1/52024-07-13

Summary

This paper addresses the challenge of solving linear inverse problems in high-resolution image synthesis using generative models based on flow matching. While these models are appealing due to their ability to compute image likelihoods directly from a learned flow, they suffer from slow log-likelihood computations that involve backpropagation through an ODE solver. This computational bottleneck can be particularly problematic for high-dimensional problems. To overcome this issue, the authors propose a novel iterative algorithm designed to approximate the maximum-a-posteriori (MAP) estimator efficiently. The key insight is to decompose the MAP objective into multiple "local MAP" objectives, leveraging Tweedie’s formula to perform gradient steps for sequential optimization. This approach allows for a more efficient approximation of the MAP estimator. The paper validates the proposed algorithm across various linear inverse problems, such as super-resolution, deblurring, inpainting, and compressed sensing. The results show that the new method outperforms existing techniques based on flow matching, making it a promising tool for high-resolution image synthesis and related tasks. The authors conclude that their approach successfully addresses the computational challenges of flow matching models, offering a robust solution for incorporating flow priors in solving linear inverse problems. They also mention discussing limitations and future work in the appendix.

Strengths

- originality: the authors present an original idea to utilize the image probabilities obtained from the learned flows in Flow Matching as priors for MAP estimation in inverse problems - quality: The authors provide a good account of Flow Matching and deep understanding of the inverse problems. Delivery of the argument is elaborated and supported with theoretical proofs as well as intuitive toy data examples and more complex tasks. Ablation studies to support the argument are provided. - clarity: fairly clear, improvements on the flow of argument would be desired. - significance: applications to real-world imaging problems such as medical imaging are very important, and the authors provide promising results in section 4.2, including Table 1 and Figure 4 and Appendix G.

Weaknesses

The the paper could benefit from more clear presentation, tying the proofs and experimental results to the claims introduced at the beginning a bit tighter.

Questions

NA

Rating

7

Confidence

1

Soundness

3

Presentation

3

Contribution

3

Limitations

NA

Reviewer hhXk5/10 · confidence 4/52024-07-14

Summary

The paper proposes efficiently recovering MAP estimates by utilizing flow-matching priors. The main proposition in the presented method is to break down the MAP objective into a sum of N local MAP objectives, facilitating a computationally feasible approach that runs in reasonable time. The results are compared to several proposed baselines in terms of distortion (PSNR/SSIM) showing consistent improvements across two datasets and multiple tasks.

Strengths

* The paper proposes a theoretically motivated approximation of MAP estimates utilizing flow-matching priors. * The paper is well written and overall well structured. * The experiments entail multiple tasks and datasets, consistently showing improved performance in distortion.

Weaknesses

* The paper argues that flow-matching priors are useful due to the ability to calculate reconstruction log-likelihood, yet this information is never used and does not appear in the experiments section * The authors focus only on distortion (and not on perceptual quality), arguing for a single (blurry) MAP estimate. Nonetheless, by this point it is well understood within the image restoration community that summarizing posteriors to a single prediction is inevitably throwing away uncertainty information which is vital for proper down-the-line decision-making in safety-critical applications. Therefore, posterior samplers (aka stochastic inverse problem solvers) are a much better solution in underdetermined inverse problems with a multitude of admissible solutions. * Nowadays posterior samplers utilizing distilled diffusion models can produce samples with a single neural function evaluation. Compared to such methods, it is unclear what might be the advantage of the presented technique besides maybe accompanying each reconstruction with its likelihood, which is again missing in the presented results.

Questions

* What information is there to be gained from the likelihood computation if all the user is going to be presented with is a single MAP estimate? * In terms of runtime, how does your method fair with other types of MAP estimates employing different generative models as priors? * What changes will it take to apply your method to non-linear inverse problems?

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

The authors were upfront about the limitations of their method and stated these in Appendix A.

Authorsrebuttal2024-08-12

Dear reviewer, Thanks a lot for your time and effort. As the discussion period approaches its deadline, we would greatly appreciate your feedback on whether our rebuttal has adequately addressed your concerns. Please reach out if you have additional questions or comments which we are happy to address. Thank you again. Best, Authors

Reviewer hhXk2024-08-12

Thanks for the clear and thorough rebuttal. I'm hesitant to change my score mainly because I do not agree with you on two points: * First, the likelihood information, if the user is to be presented with a single reconstruction, is practically useless as ultimately what matters is the likelihood ratio between different possibilities. * Second, the results are extra blurry. As evident in the FID numbers from the rebuttal PDF, this method is hardly competitive in perceptual quality. I do recognize the advantage of this method over the concurrent work of D-Flow. Nonetheless, in its current form, I think it has limited applicability and therefore I maintain my score.

Authorsrebuttal2024-08-13

We greatly appreciate your time and recognition of our paper's strengths compared to the concurrent work, D-Flow. We would like to provide a last-minute clarification on the two points you mentioned, which we hope will address some of your concerns: - **Regarding likelihood information:** We appreciate that giving the user likelihood estimates in a more systematic way could be useful, but we want to highlight that the main focus of our work is using a flow-based prior's likelihood to help obtain high-quality reconstruction estimates for inverse problems in a computationally efficient way. As compared to baselines such as D-Flow, we can do this in a computationally efficient manner based on our theoretically-motivated MAP approximation. Moreover, focusing on reconstruction enables a fair comparison with flow-based baselines such as D-Flow and OT-ODE, as well as diffusion-based methods like RED-Diff and $\Pi$GDM. Of note is that likelihood information has also not been provided in the experiments of the four baselines. - **Perceptual quality:** We would like to highlight that our method demonstrates strong competitiveness in perceptual quality, as evidenced by the fact that our method outperforms $\Pi$GDM in terms of FID in super-resolution and inpainting (random), surpasses OT-ODE in both FID and LPIPS in inpainting (random) and inpainting (box), and exceeds RED-Diff in both FID and LPIPS across all tasks except Gaussian Deblurring. As shown in the Figure 1 in the rebuttal pdf, our method generates clear images even in the highly noisy setting and is strong in terms of PSNR and SSIM. Overall, our method achieves consistently improved performance in reconstruction while maintaining competitive perceptual quality compared to the baselines.

Reviewer hhXk2024-08-13

Thank you for these additional clarifications. After reading the reviews of others, re-reading the paper, and looking again into the rebuttal PDF, I would like to raise my score to 5 to enhance your chances of getting accepted. Best of luck!

Authorsrebuttal2024-08-13

We greatly appreciate the time and effort you've dedicated to our work. Thank you for your support and thought-provoking suggestions!

Authorsrebuttal2024-08-10

Kind reminder

Dear all reviewers, Thank you for all the insightful and constructive comments and feedback, which greatly helped us to improve the paper quality. This is a kind reminder that the discussion period is closing in a few days. Feel free to check our response to your questions and tell us if you have any further concerns. Best, Authors

Reviewer sT2G2024-08-11

Response to the Rebuttal

Thank the authors for the rebuttal. I strongly recommend the authors to characterize more about the non-asymptotic case of c(N), given that the implementation is indeed not exactly the same as the assumption made in asymptotic case. The authors should explicitly clarify this in the paper, especially in the Theorem part. I have no question about the comparison with the traditional methods and the experiment settings. I appreciate the authors for their hard work on the additional experiments and clarification. After careful consideration, I keep the original rating, but increase the contribution rating and my confidence.

Authorsrebuttal2024-08-14

A summary of the discussions

Dear AC, Thank you for your time and effort in evaluating our work. We would like to provide a brief summary of the discussions between the reviewers and us to facilitate a better evaluation. Overall, our responses to each reviewer's questions received positive feedback. - **Reviewer hhXk:** After acknowledging our method's strength over the recent concurrent work, D-Flow, the reviewer expressed major concerns regarding likelihood information and perceptual quality in their second-to-last comment. We emphasized that our method focuses on computationally efficient, high-quality reconstructions with highly competitive perceptual performance compared to other baselines, as evidenced by our rebuttal PDF. Following our response, the reviewer stated, *raised my score to enhance your chances of getting accepted.* - **Reviewer sT2G:** The reviewer’s main concerns were about the comparison with traditional methods, more challenging experiment settings, and the need for clearer specifications in our theorem statement. After our response, the reviewer had no further questions on the experiments required and stated, *increase the contribution rating and my confidence.* - **Reviewer qjx2:** The reviewer acknowledged that we *have addressed my concerns by providing empirical evidence for the requested cases*, such as severely ill-posed inverse problems in highly noisy settings, with more baselines compared and perceptual metrics reported. Consequently, the reviewer chose to *raise the score to 5*. We hope this summary helps you better evaluate the discussion process. Best, Authors

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC