Flow Priors for Linear Inverse Problems via Iterative Corrupted Trajectory Matching

We thank you for your helpful review and provide the additional experiments which we hope address your concerns:

More baselines: RED-Diff, $\Pi$GDM, and D-Flow See Tab. 2 in the attached pdf above. In addition to the flow-based baselines, we have included two representative diffusion-based baselines: 1) RED-Diff [1], a variational Bayes-based method; and 2) $\Pi$GDM [2], an advanced MCMC-based method. Our method demonstrates competitive performance with both diffusion-based baselines and other approaches in terms of recovery fidelity.

Thank you for bringing D-Flow [3], recently accepted to ICML 2024, to our attention. We did not compare with D-Flow as we consider it concurrent work and it has not yet released its official implementation. Notably, D-Flow also includes [1] and [2] as diffusion-based baselines. Implementation details are included at the end of the [Author Rebuttal] above.

Runtime comparison with D-Flow As shown in the Table 3 in the appendix of our paper, our method requires 1.6 min for each image, OT-ODE [2] 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.

Highly noisy setting See Tab. 2 of the attached pdf above. We present the results for $\sigma_y = 0.2$ along with a qualitative comparison at the end of the document. We adopted the same hyperparameter selection strategy as described in the reference paper. With a fixed step size of $\eta = 10^{-2}$, we observed that decreasing the guidance scale is beneficial in very noisy conditions. Specifically, we chose $\lambda = 10^{2}$ for super-resolution, random inpainting, and Gaussian deblurring, while $\lambda = 10^{1}$ was selected for box inpainting. Our method outperforms other flow-based baselines, demonstrating its robustness in handling very noisy scenarios.

Severely ill-posed problems See Tab. 1 of the attached pdf above. We consider increasing the level of ill-posedness in the compressed sensing MRI experiments, as this is a more challenging inverse problem. We set the compression rate to be $\nu=1/10$, i.e. only 10% of the output signal has been observed, which relates to inpainting with 90% of pixels missing you mentioned. Our method outperforms the classical recovery algorithms and other baselines in all settings, demonstrating our method's capability to handle challenging scenarios and the advantages of utilizing modern generative models as priors. Implementation details are included at the end of the [Author Rebuttal] above.

LPIPS FID scores They are provided in Tab. 3 of the attached pdf above. Note that in papers oriented by posterior sampling, more emphasis is put on perceptual quality metrics, such as FID and LPIPS. A large number of images (usually 1k) is thus required due to FID's Gaussian assumptions on the distribution. However, our paper focuses mainly on metrics of recovery fidelity such as PSNR and SSIM, which can be calculated by individual image pairs (the same applies to LPIPS). We find both metrics remain stable after the number of test images reach 100. Note that our method is also competitive with other baselines in terms of perceptual quality as evidenced by our LPIPS scores. The FID score calculated by 100 images is quite noisy and just for reference.

We hope our additional experiments on severely ill-posed inverse problems, in highly noisy setting, with more baselines compared and perceptual metrics reported can address some of your concerns. Finally, we want to highlight our theoretical contribution to inverse problems with flow-based models as priors, where backpropagation through an ODE solver (usually requiring 100 or more NFEs) is avoided, and asymptotic convergence is proven, as demonstrated in Theorem 1. We kindly request a re-evaluation of the rating if some of your concerns have been resolved.

[1] A Variational Perspective on Solving Inverse Problems with Diffusion Models. Mardani, Morteza and Song, Jiaming and Kautz, Jan and Vahdat, Arash. ICLR 2024.

[2] Pseudoinverse-guided diffusion models for inverse problems. Song, Jiaming and Vahdat, Arash and Mardani, Morteza and Kautz, Jan. ICLR 2023.

[3] 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)

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

We sincerely thank you for your helpful and detailed feedback, which significantly helped us improve the paper quality.

Dependence on $N$ in Theorem 1 We appreciate your concern regarding our theory. While we agree that it could be beneficial to obtain a non-asymptotic error bound in terms of $N$, the goal of this Theorem is to provide intuition for why our local MAP objectives are a reasonable approximation to the global MAP problem. The asymptotic nature of the Theorem simplifies its presentation, while still providing support for why our approach works. The terms that actually govern how the error scales with $N$ are the $\hat{\mathcal{J}}_i$ terms. Obtaining non-asymptotic error bounds in $N$ is possible, but this would require further technical assumptions on the flow model to control the approximation error via Riemannian sums.

The $c(N)$ term is given by $c(N) = \sum_{i=1}^N \gamma_i c_i - \log p(y) = \sum_{i=1}^N (\frac{1}{2})^{N-i+1} m\log (\frac{i}{N} ) - \log p(y)$. This quantity converges to a constant as $N \rightarrow \infty$. We promise to include the explicit expression of $c(N)$ into Theorem 1 in the revision.

Classical recovery algorithms: Wavelet and TV priors Thank you for this suggestion. See Tab. 1 in the attached pdf above. We have added results for two classical recovery algorithms, Wavelet and TV priors. Based on the first two rows, the TV prior consistently surpasses the Wavelet prior. We surmise that this is because the TV prior is more effective for images with clear boundaries and homogeneous regions, such as those in the HCP T2w dataset. We find that our method outperforms the classical recovery algorithms in all settings. Implementation details are included at the end of the [Author Rebuttal] above.

More challenging setting: smaller compression rate $\nu=1/10$ Thank you for this suggestion. See Tab. 1 in the attached pdf above. We conduct an additional experiment using the compression ratio $\nu = 1/10$, i.e only 10% of the output signal has been observed. Based on the two rightmost columns of Tab.1, we see that our method consistently outperforms all the other baselines in terms of PSNR and SSIM.

More baselines: RED-Diff and $\Pi$GDM See Tab. 2 in the attached pdf above. In addition to the flow-based baselines, we have included two representative diffusion-based baselines: 1) RED-Diff [1], a variational Bayes-based method; and 2) $\Pi$GDM [2], an advanced MCMC-based method. Our method demonstrates competitive performance with both diffusion-based baselines and other approaches in terms of recovery fidelity. Additionally, one reviewer mentioned a concurrent work, D-Flow [3], recently accepted to ICML 2024, which has not yet released its official code. Notably, D-Flow also includes [1] and [2] as diffusion-based baselines. Implementation details are included at the end of the [Author Rebuttal] above.

Hyperparameter tuning for $\lambda, \eta, K$ In our implementation, we found that with the Adam optimizer, the choice of step size $\eta$ and guidance weight $\lambda$ remains consistent. Practically, for a new task, we recommend starting by fixing the iteration number $K$ at 1. An initial step size of $\eta=10^{-2}$ is likely to perform well, as this value was effective across all tasks, from natural images to medical applications, after a comprehensive grid search (refer to Fig 5 in the paper). The value of $\lambda$ should then be determined by a grid search within the set {$10^2, 10^3, ..., 10^7$}. After establishing $\lambda$, we suggest incrementally increasing $K$ to observe if performance improves. If performance continues to increase, keep raising $K$ until the metrics plateau; if there is no improvement, $K=1$ is the optimal choice.

We hope that our additional experiments in more challenging setting, with classical recovery algorithms and representative diffusion-based methods compared can address some of your concerns. We kindly request a re-evaluation of the rating if some of your concerns have been resolved.

[1] A Variational Perspective on Solving Inverse Problems with Diffusion Models. Mardani, Morteza and Song, Jiaming and Kautz, Jan and Vahdat, Arash. ICLR 2024.

[2] Pseudoinverse-guided diffusion models for inverse problems. Song, Jiaming and Vahdat, Arash and Mardani, Morteza and Kautz, Jan. ICLR 2023.

[3] 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)

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)

We appreciate your recognition of our original contributions and will work on enhancing the clarity and coherence of our presentation. In the attached pdf above, more experiments have been added to show our method's capability to handle more challenging conditions ($\nu=1/10$, Tab. 1), highly noisy settings (Tab. 2), strong performance over classical signal recovery algorithms (Tab. 1) brought by utilizing generative models as priors, and its competitiveness with representative diffusion-based methods (Tab. 2).

We thank the reviewer for their detailed feedback and comments. We would like to address your concerns regarding our paper's motivation and contribution to the inverse problems community. We will respond to each concern below:

The use of log-likelihood and flow prior

First, we would like to note that our algorithm does use the likelihood under a flow prior. In particular, our algorithm optimizes a sequence of objectives that locally approximate the log-likelihood for each timestep $t$. In the lines 7 and 12 of Algo. 1 of the paper, the terms $\frac{1}{2}||x_t||$, $\log p(x_t)$, and the trace approximation are derived from the flow prior's log likelihood. We can obtain the gradient of $\log p(x_t)$ using Eq (11). We experimentally demonstrate that these prior terms are important in obtaining quality reconstructions in Fig. 2 by comparing our results to the results of ours without the local prior terms. Hence the flow prior and log-likelihood is heavily used in our algorithm, and are crucial to help us obtain good performance.

Second, with flow-based models, we are able to calculate $\log p(x)$ by the instantaneous change-of-variable formula (Eq. 8 in the paper). We choose not to present the log-likelihoods and mainly present reconstruction metrics, such as PSNR and SSIM, to properly compare our method with flow-based baselines such as D-Flow [1] and OT-ODE [2], and diffusion-based methods like RED-Diff [3] and $\Pi$GDM [4]. We fully agree that reporting log-likelihoods could help users gain more information of the reconstructed images and we promise to do so in the revision.

On MAP Estimation and Uncertainty Quantification Thank you for raising this issue. We acknowledge that posterior sampling methods have the advantage that they can provide multiple reconstructions to visualize uncertainty. However, most methods that we compare to, such as our baselines of [1] [2] [3] and [4], only provide one estimate for each single image and do not focus on uncertainty quantification. We would like to highlight, however, that it is possible to equip MAP with uncertainty quantification, which we hope to explore in future work. For example, given an MAP solution $\hat x_{MAP}$, one can compute a Laplace approximation $p(x_{MAP}|y) \sim N(\hat x_{MAP}, H^{-1})$ where $H$ is the Hessian matrix of $\log p(x|y)$ at $\hat x_{MAP}$. This method has its pros and cons relative to posterior sampling methods. Here, computing this Hessian can be challenging, especially for flow-based log priors, but once calculated, sampling is straightforward. Similarly, posterior sampling methods require many iterations to generate many quality samples, which can be time-consuming. We believe investigating the benefits of this approach constitutes an interesting direction for future work.

On Distilled Models While distilled models offer a promising way to speed up inversion with diffusion models, to our knowledge there have been few works that have successfully demonstrated their strong performance. Moreover, one-step distilled diffusion models behave similarly to push-forward generators, such as GANs. Even with distilled diffusion models, both posterior sampling and MAP estimation require many optimization steps to achieve satisfactory reconstructions, as shown in classical papers using GAN priors. See, for example, Sec 6.1.1 in CSGM [5], where they optimize the latent space of the GAN for 1000 steps to reconstruct an MNIST digit.

Runtime Comparison As shown in the Table 3 in the appendix of our paper, our method requires 1.6 min for each image, OT-ODE [2] 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.

Finally, we want to highlight our theoretical contribution to inverse problems with flow-based models as priors, where backpropagation through an ODE solver (usually requiring 100 or more NFEs) is avoided, and asymptotic convergence is proven, as demonstrated in Theorem 1. Practically, our ICTM algorithm consistently achieves strong performance in distortion across multiple tasks and datasets, from natural images to medical application, as evidenced by Tabs. 1 and 2 in the attached pdf above. We kindly request a re-evaluation of the rating if some of your concerns have been resolved.

[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)

[2] Training-free linear image inverses via flows. Ashwini Pokle, Matthew J. Muckley, Ricky T. Q. Chen, Brian Karrer. TMLR 2024.

[3] A Variational Perspective on Solving Inverse Problems with Diffusion Models. Mardani, Morteza and Song, Jiaming and Kautz, Jan and Vahdat, Arash. ICLR 2024.

[4] Pseudoinverse-guided diffusion models for inverse problems. Song, Jiaming and Vahdat, Arash and Mardani, Morteza and Kautz, Jan. ICLR 2023.

[5] Compressed Sensing using Generative Models. Ashish Bora, Ajil Jalal, Eric Price, and Alexandros G. Dimakis. ICML 2017.

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

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.