Flow Priors for Linear Inverse Problems via  Iterative Corrupted Trajectory Matching
Authors: Yasi Zhang
Peiyu Yu
Yaxuan Zhu
Yingshan Chang
Feng Gao
Ying Nian Wu
Oscar Leong
Venue: NeurIPS 2024 poster (2024)
Keywords: Flow-based models
Inverse problems
Neural ODE
Primary area: generative_models

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 \url{https://github.com/YasminZhang/ICTM}.

Flow Priors for Linear Inverse Problems via Iterative Corrupted Trajectory Matching
Yasi Zhang
Peiyu Yu UCLA yupeiyu980g. ucla..0..ucledee ucla.edu
Yaxuan Zhu UCLA yaxuanzhudg. ecla.edu
UCLA
yascinzhang@ucla.edi
Yingshan Chang
Feng Gao"
Ying Nian Wu UCLA
CMU
Amazon
yingshac@andrew.com cmu. edu
fenggo@amazon.com com
ywuêstat ucla. edu
Oscar Leong
UCLA
oleong@stat. ucla.
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-variable formula, one can directly compute image likelihoods from learned flow, making them enticing candidates as priors for downstream tasks such as inverse problems. In particular, a natural approach would be incorporate such image probabilities in estimation problem. major obstacle, however, lies in the slow computation the logeloo,,,,od, requires backpropagating through ODE solver, which can be prohibitively slow for high-dimensional problems. In this work, we propose iterative algorithm to approximate the MAP estimator efficiently to solve variety of linear inverse problems. Our algorithm is mathematically justified by the observation that the MAP objective can be approximated by sum of N "local MAP" objectives, where N is the number of function evaluations. By leveraging Tweedie's 's formula, 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 available at https: / githuu........ wtt.///.o./////.eee WasminZhang/IITT
(MAP)
show that can perform gradient steps sequentially
Introduction
Linear inverse problems are ubiquitous across many imaging domains, pervading areas such as astronomy [41, 23], medical imaging [38, 49], and seismology [35, 39]. In these problems the goal is reconstruct an unknown image æ, R from observed measurements y 1 of the form:
(I) where A Rd + RM with < is linear operator that degrades the clean image F., and the additive noise drawn from known distribution. In this work, we assume the noise follows (0 I). Due the under-constrained nature of such problems, they are typically ill-posed, i.e.,
A(x.) + noise,
"This is not related to author's position Amazon.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

there are an infinite number of undesirable images that fit to the observed measurements. Hence, one requires further structural information about the underlying images, which constitutes our prior.
With the advent of large generative models [27, 17, 48, 40, 8, (555558 59), there has been a surge of interest in exploiting generative models as priors to solve inverse problems. Given pretrained generator to sample from a distribution or grant access to image probabilities, one can solve variety of inverse problems in task- or forward model-agnostic fashion, without the need for large-scale supervision [36]. This has been successfully done for a variety of models, including implicit generators such as Generative Adversarial Networks (GANs) and Variational Autoencoders (VAEs) [4, 34], invertible generators such as Normalizing Flows [1, 54], and more recently Diffusion models [10, 43, 60].
A recent paradigm in generative modeling [48, 55, 25, 57], based on the concept of flow matching [29, 28), has made significant strides in scaling ODE-based generators to high-resolution images. Flow matching models map simple base distribution, such a Gaussian, to a complex, high- dimensional data distribution by defining a flow field that represents the transformation between these distributions. These generative models have demonstrated scalability high dimensions, forming the backbone of several state-of-the-an generative models [30, 13, 56]. Moreover, flow matching models follow straighter and more direct probability paths compared diffusion models, allowing for more efficient and faster sampling [28, 29, 13J. Additionally, due to their invertibility, flow matching models provide direct access to image likelihoods through the instantaneous change-of- of-variables formula [9. 18]. Given these advantages and the relatively recent application of these models to inverse problems [37. 2]. we investigate their use as image priors in this work.
Leveraging knowledge about the corruption process p(y (r) and a natural image prior p(x), the Bayesian approach suggests analyzing the image reconstruction posterior p(xly) xx p(y]=))(()) to solve the inverse problem. A proven and effective method based on this approach is maximum-a- posteriori (MAP) estimation [6. 19], which maximizes the posterior identify the image most likely to match the observed measurements:
argmin ((xy)) argmin log p(ylar) log p(x).
(2)
=ER"
zER"
MAP estimation provides single, most probable point estimate of the posterior distribution, making simple and interpretable. This deterministic approach ensures consistency and reproducibility, which are essential in applications requiring reliable outcomes, particularly in compressed sensing tasks such as Computed Tomography (CT) [7] and Magnetic Resonance Imaging (MRI) [52]. While posterior sampling methods can offer diverse reconstructions to quantify uncertainty, they can be prohibitively slow in high-dimensions [5]. Hence, in this work, we propose to integrate flow priors to solve linear inverse problems by MAP estimation.
significant challenge in employing flow priors for MAP estimation lies in the slow computation of the image probabilities, as it requires backpropagating through an ODE solver [47, 16, 15]. In this work, we show how one can address this challenge via Iterative Corrupted Trajectory Matching (ICTM), a novel algorithm to approximate the MAP solution in a computaionally efficient manner. In particular, we show how one can approximately find an MAP solution by sequentially optimizing a novel simpler, auxillary objective that approximates the true MAP objective in the limit of infinite function evaluations. For finite evaluations, we demonstrate that this approximation is sufficient to optimize by showcasing strong empirical performance for flow priors across a variety of linear inverse problems. We summarize our contributions as follows:
1.
We propose ICTM, an algorithm to approximate the MAP solution to a variety of linear inverse problems using a flow prior. This algorithm optimizes an auxillary objective that partitions the flow model's trajectory into N "local MAP" objectives, where N is the number of function evaluations (NFEs). By leveraging Tweedie's formula, we show that we can perform gradient steps to sequentially optimize these objectives.
2. Theoretically, we demonstrate that the auxillary objective converges to the true MAP objective as the NFEs goes to infinity. We validate the correctness of our algorithm in finding the MAP solution on denoising problem.
3. We demonstrate the utility of ICTM on a wide variety of linear inverse problems on both natural and scientific image datasets, with problems including denoising, inpainting, super- resolution, deblurring, and compressed sensing. Extensive results show that ICTM is both computationaly efficient and obtains high-quality reconstructions, outperforming other reconstruction algorithms based on flow priors.

---

2 Background
Notation We follow the convention for flow-based models, where Gaussian noise is sampled at timestep 0, and the clean image corresponds to timestep 1. Note that this is the opposite of diffusion models. Fort E 10, 1], we denote ze (To) as the point at time t whose initial condition is zu- this work, we use z and 21 interchangeably. i.e, 21 (xo) = r(2o).
2.1 Flow-Based Models
We consider generative models that map samples z6 from noise distribution p(zo), e.g., Gaussian, to samples Iı of a data distribution p(æ1) using an ordinary differential equation (ODE);
dx, = (z4. t) dt.
where the velocity field is a 0-parameterized neural network, e.g., using a UNet [28, 29, 42] or Transformer [13, 51] architecture. Generative models based on flow matching [28, 29] can be seen as a simulation-fiee approach to learning the velocity field. This approach involves pre-determining paths that the ODE should follow by specifying the interpolation curve rather than relying on the MLE algorithm to implicitly discover them [9]. To construct such a path, which is not necessarily Markovian, one can define differentiable nonlinear interpolation between and 21:
It = 04 + 3,ro, To ~ NN0, I):
44)
where both a and 3. are differentiable functions with respect to satisfying ag = 0. 80 = 1, and Qi 1, 31 = 0. This ensures that is transported from a standard Gaussian distribution to the natural image manifold from time to time 1. contrast, the diffusion process [48, 45, 20) induces non-differentiable trajectory due the diffusion term in the SDE formulation.
The idea behind flow matching the velocity field at each timestep. To achieve this, we can train the neural network by minimizing an L2 loss between the sampled velocity and the one predicted by the neural network:
utilize the power of deep neural networks efficiently predict
(()) = Pe.p((.). p(po)) [[ ( 1) - (ó,31 + 3:
(5)
We denote the optimal (not necessarily unique) solution arg ming L(0) as ô. The optimal velocity field by can be derived in closed form and is the expected velocity at state Ny:
ug(#,, t) =
3,ro xz].
(6)
For convenience, in the following text, we use uo to refer the optimal vg- In the rest of the paper, we assume that the flow uy and its parameters are pretrained on dataset of interest and fixed. We are then interested leveraging its utility as a prior to solve inverse problems.
2.2 Probability Computation for Flow Priors
Denote the probability of x; in Eq. (3) as p(xz) dependent on time. Assuming that va is uniformly Lipschitz continuous in 3n and continuous in I, the change in log probability also follows a differential equation [9, 18]:
log p(xx) - ( uo (2 (2), Zx
Ø1
One can additionally obtain the likelihood of the trajectory via integrating Eq. (7) across time
8) ( ds, 0 O
Or (2
log p(z.) = log p(x,)
tr
(8)
3 Method
In this work, we aim to solve the MAP estimation problem in Eq. (2) where p(z) is given by pretrained flow prior. We first discuss Section 3.1 how the MAP problem could, in principle, be solved via a latent-space optimization problem. As will see, this problem is challenging solve

---

computationally due to the need to backpropagate through an ODE solver. To overcome this, we show in Section 3.2 that the ideal MAP problem can be approximated by weighted sum of "local MAP" optimization problems, which operates by partitioning the flow's trajectory to reconstructed solution. We then introduce our ICTM algorithm to sequentially optimize this auxiliary objective. Finally, in Section 3.3, we experimentally validate that our algorithm finds a solution that is faithful to the MAP estimate in a simplified setting where the globally optimal MAP solution is known.
3.1 Flow-Based MAP
Given pretrained flow prior, one can compute the log-likelihood of x generated from an initial noise sample 26 via Eq. (8). Hence, to find the MAP estimate, could equivalently optimize the initial point of the trajectory zo and return xj (2o) where zo is found by solving
! tr prior
min #o #oER" 20 lly A(21( |
Ox vo (xc: dt.
(9)
data likellhood
where T2 := (20) denotes the intermediate state generated from zo. Intuitively, this loss encourages finding an initial point To such that the reconstruction I1 == x1 (2a) fits the observed measurements, but is also likely to be generated by the flow.
In practice, 21 and the prior term be approximated by ODE solver. The trajectory of Ie Ho+ d0 (z t)dt can be approximated by an ODE sampler, i.c. ODESolve(zo- t, vy), where zo is the initial point, and the second and third arguments represent the starting time and the ending time, respectively. For example, with an Euler sampler, we iterate over Z++ I, (o At where At = 1/N and N is the predetermined NFEs. After acquiring the optimal 2o by optimizing the Eq. (9), we obtain the MAP solution 21 by using ODESolve($o, 0, vg) again.
3.2 Flow-Based MAP Approximation
The global flow-based MAP objective Eq. (9) is tractable for low-dimensional problems. The challenge for high-dimensional problems, however, is that optimizing Eq. (9) is simulation-based. and thus each update iteration requires full forward and backward propagation through an ODE solver, resulting in issues regarding memory inefficiency and time, making it hard optimize [9, 15, 16, 47].
As way to address this, we prove a result Theorem I that shows that the MAP objective can be approximated by weighted sum of N local posterior objectives. These objectives are "local" in the sense that they mainly depend on likelihoods and probabilities of intermediate trajectories de and + eg (Iz, At for t 0, At..... N At where At := 1/N. Given an initial noise input Zo. each local posterior objective depends on a non- Markovaan auxiliary path y/ = a,y + 3(A(1o) by connecting the points between y and Azg. We prove this result for straight paths aę t and 3. = =1= for simplicity, but other interpolation paths can be used. The proof is in Section A.2.
Theorem 1. For N 1, 7 == (2) N-)+ and At 1/N. Suppose y A(xx)) + where v. xn (30) with 3o being the solution Eq. (9), 6~ of,, and xa exactly follows the straight path tx (1 t) for any timestep t [0, 1]. Suppose the velocity field uo R x R + R satisfies supzeR=_xE(o. tr 2 vo(2,8) C1 for some universal constant C1. Then, there exists constant N that does not depend on To such that
N Mr log logp(z(zo)lv) - Ench -c(N) =0,
lim
N
7=1
where J. = log (ogp((((11))))) - tr (800(tn- (11A)t)) At + log p(yisi (.a))).
This result shows that the true MAP objective evaluated at the optimal solution can be approximated by a weighted sum of objectives that depend locally at a time for the trajectory {xt | E [o, 1]]. The intuition regarding J, arises from the fact that J, s J. where J, is the local posterior distribution
J. = log M Dise ( + log p(ziat).
This is given c(N) ΣN 7ec log) Please the proof of Theorem in Appendix A.2.

---

Optimizing each of these local posterior distributions in sequential fashion captures the fact that we would like each intermediate point in our trajectory 2, to be likely and fit to our measurements, ideally resulting final reconstruction #1 that satisfies this as well. The benefit of J., as we will show in the sequel, is that it is efficient optimize.
Discussion assumptions: We assume that the trajectory {xx )s exactly follows the predefined interpolation path (0. + h- In Section B of the appendix, we analyze this assumption and show that we can bound the deviation from the predefined interpolation path to the learned path via path compliance measure. Moreover, impose a regularity assumption on the velocity field vg. effectively requiring a uniform bound on the spectrum of the Jacobian of Ug. This can be easily satisfied with neural networks using Lipschitz continuous and differentiable activation functions.
As we see in Theorem 1, one can approximate the true MAP objective via a sum of local objectives of the form
tr (Ouo(zes 1) (i 1)2()
+ +log 1ogp(((-1))1)) J = log (og
At.
(10)
Or
local data likelihood
local prior
At first glance, J, still appears challenging to optimize, but there are additional insights we can exploit for computation. We discuss each term in T, below.
measurement manifold
Local data likelihood: The intuition behind ICTM is that we aim to match corrupted trajectory {we} )t with an aux- iliary path {y: he specified by an interpolation between our measurements y and A(za) for each timestep t, defined by 3/ := 3i The corrupted trajectory u, :- A(x;) follows the corrupted flow ODE du, = (1,, ())dl.
dur Alsuten fildt
Me
To optimize the above "local MAP" objectives, we must understand the distribution of p(v) p())/z)). Generally speak- ing, this distribution is intractable. However, by assuming exact compliance of the trajectory generated by flow to the predefined interpolation path (as done in Theorem 1), can show that ya 2 [M,, afoz). This is proven in Lemma 3 in the appendix. While exact compliance of the trajectory may not hold for learned flow matching models, we show empirically that making this assumption leads to strong performance in practice. We further analyze this notion of compliance in Section B of the appendix.
natural image manifoid
Figure 1: Illustration of the idea of ICTM. The corrupted trajectory u, A(x;) follows the corrupted flow ODE
du, A(vo (1,, dt.
Local prior: The approximation in Eq. (10) addresses of the main concerns of MAP in that the intensive integral computation is circumvented with a simpler Riemannian sum. This approximation holds for small time increments At: p tr au vo (2). ds ss tr & (Z t)) At. Note that one can additionally improve the efficiency of this term by employing a Hutchinson Skiilling estimate first glance, it appears we have simply
[44, 21] for the trace of the Jacobian matrix. However, shifted the problem to the computation of the prior at timestep (i 1) At. Fortunately, it is possible to derive formula for the gradient of log p(x:) for all timesteps t j0, using Tweedie's 's formula [12]. This allows us optimize each objective J. using gradient-based optimizers. The following result gives precise characterization of V., log p(z,), proven in Section A.1.
Proposition 1. Let dç = au/ Be denote the signal-to to-noise ratio. The relationship between the score function V., log p(xx) and the velocity field vo (x., t) is given by:
1 B
(
(dlogg di
)
log Br Zt
Vz, logp(z))
(o(
(1I)
dl
In summary, we have derived an efficient approximation to the MAP objective. For our algorithm, we ., N NA, fitting our current iterate
iteratively optimize each term J. sequentially for each = 0, At,...,
Iz to induce an increment Z such that A(z(+4e) fits to our auxiliary corrupted path y++Ar while
5

---

Algorithm Iterative Corrupted Trajectory Matching (ICTM) with Buler Sampler
Input:
measurement y, matrix A, pretrained how-based model 6, NFEs N, interpolation coefficients fas) and {Be h. step size 7, guidance weight A, and iteration number K
Output: recovered clean image æ1
Initialize (O. I), 26 + e, 0, At 1/N
2: Generate an auxiliary path = a.y + B. (Aro) for s E (0,1)
3: while
4: I++A: < Zt + (1 t)
5: if then
6: for k 1, K do
7: St f- Zt Vv (((A(m==)s(z1)) B1t0ell + 100e|1 tr
(ovo(ze))) Ox
At)
8:
end for
9:
else
for 1,... K do
11:
# use Eq. (11) to obtain the gradient of log p(rz)
12:
It t xt V.. (()=)) vetaell p(zp( + tr
(Du((eed)
At
Ox
end for end if
14:
T++At Iz (0 t)
16:
t 1-t+ + At
end while
18: return 21
0.35
0.2.
0.30
MAP
0.05
0.15
Ours
0.0
0.6
0.00
Difference Value
NFP
(a) Qualitative Results
Histogram Differences
(c) MSE VS NFES
Figure 2: Results 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 NEEs varies.
being likely under local prior. We call this approach Iterative Corrupted Trajectory Matching (ICTM). Our algorithm is summarized in Algo. In lines and 12, instead of directly optimizing the local data likelihood, choose À as new hyper-parameter tune. We find constant À works well in practice.
3.3 Toy Example Validation
We experimentally validate that the reconstruction found via ICTM is close to the optimal MAP solution in a simplified denoising problem where the MAP solution can be obtained closed-form. Specifically, fit Gaussian distribution M(,, Σ) using 1,000 samples from the FFHQ dataset. Consider a denoising problem
x + where x ~ N N(4, Σ) and (N0, o2 this case, the analytical solution to the MAP estimation problem (Eq, (2)) is = (E-'+o, - uto, 2v). set og = 0.1. Then, we train a flow-based model on 10,000 samples from the true Gaussian distribution and showcase the deviation of our reconstruction found via ICTM to the closed-form MAP solution v, in Fig. 2. see that ICTM can obtain faithful estimate of the MAP solution across many samples.

---

Table I: Quantitative comparison results in terms of PSNR and SSIM on the CelebA-HQ dataset. Our algorithm surpasses all other baselines across tasks. The best values are highlighted in blue and the second-best are underlined.
Super-Resolution
Inpainting(random)
Gaussian Deblurring
Inpainting(bxx)
Method
PSNR
SSIM
PSNR
SSIM
PSNR
SSIM
PSNR
SSIM
OT-ODE
27.46
0.775
28.57
0.838
26.28
0.727
19.80
0.795
DPS-ODE
27.85
0.791
29.57
0.872
25.97
0.704
23.59
0.758
RED-Diff
27.20
0.760
25.13
0.711
27.23
0.765
17.50
0.651
HIGMM
28.33
0.803
29.98
0.858
24.30
0.583
24.10
0.853
Ours (w/o prior)
26.06
0.724
29.01
0.835
25.13
0.676
22.42
0.803
Ours
27.91
0.805
30.65
0.894
26.54
0.760
24.34
0.866
OT-ODE
DPS-ODE
Ours
Groundinath
Inout
QT-CDE
DPS-ODE
Ours
Groundruth
Super-Pesotution
(t) Inpaintinng (randomt)
(cI Gauzsian Dablurring
(c) Inpainting (oox)
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.
Experiments
In our experimental setting, we use optimal transport interpolation coefficients, i.e. Oe = / and 3 11 We test our algorithm on both natural and medical imaging datasets. For natural images, we utilize the pretrained checkpoint from the official Rectified Flow repository and evaluate our approach on the CelcbA-HQ dataset [31, 24]. We address four common lincar inverse problems: super-resolution, impainting with a random mask, Gaussian deblurring, and inpainting with a box mask. For the medical application,
train flow-based model from scratch the Human Connectome Project (HCP) dataset [50] and test our algorithm specifically for compressed sensing at different compression rates. Our algorithm focuses the reconstruction faithfulness of generated images, therefore employing PSNR and SSIM [53] as evaluation metrics.
Baselines We compare our method with five baselines. 1) OT-ODE [37]. To our knowledge, this the only baseline that applies flow-based models inverse problems. They incorporate prior gradient correction at each sampling step based on conditional Optimal Transport (OT) paths. For fair comparison, we follow their implementation of Algorithm 1, providing detailed ablations on initialization time in Appendix E.3. 2) DPS-ODE. Inspired by DPS [10], we replace the velocity field with conditional one, i.e., v( ly) v(x,) + V. log plyl (zi)), where & is hyperparameter to tune. Following the hyperparameter instruction in DPS, we provide detailed ablations on Çe Appendix E.3. 3) Ours without local prior. To examine the local prior term's effectiveness in our optimization algorithm. we drop the local prior term as defined in Eq. (10) in our algorithm. In the experiments with natural images, in addition the flow-based baselines, we have included two representative diffusion- method: and 5) JIGDM [46], an advanced MCMC-based method. We also note one concurrent work, D-Flow [2], which formulates the MAP as a constrained optimization problem in their Eq. 9, As documented in their Sec. 3.4, takes 5-10 minutes to recover each image. This because each of its optimization step requires backpropagation through an ODE solver to compute the full log-likelihona
baselines: 4) RED-Diff [33]. a variational Bayes based
Phttps://gitaaii http::::/g.ttb..ooo//oobtttaabba www.gnobbtab//hetiiieelee

---

Table 2: Results of compressed sensing with varying compression rate /' on the HCP T2w dataset. Note that compressed sensing is more challenging due the complexity of the forward operator, as evidenced by the poor performance of OT-ODE, which assumes a Gaussian distribution of measurement y given F/. The best values are highlighted in blue.
1/2
1/4
1/10
Method
PSNR
SSIM
PSNR
SSIM
PSNR
SSIM
Wavelet Prior
18.02 = 1,38
0.495 + 0.02
11.99 = 1.34
0.230 0.02
7.37 = 1.85
0.090 0.02
TV Prior
25.36 2.79
0.657 0.04
2.36
0.496 = 0.03
3.04
0.04
OT-ODE
18.71 = 1.02
0.422 + 0.17
18.16 1.06
0.271
0.07
12.21
1.43
0.096 0.04
DPS-ODE
31.06 3.91
0.765 0.08
25.01 1.87
0.608 0.08
22.06
1.66
0.479
0.09
Ours
32.72 1.53
0.878 0.05
27.03 1.77
0.733 0.04
24.03 1.23
0.503 0.04
OT-ODE
DPS-ODE
Ours
GT
OT-ODE
DPS-ODE
Ours
GT
A
A
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.
In contrast, our method is significantly faster (approximately 1.6 minutes per image) due to our principled local MAP approximation, as demonstrated in Appendix D.
4.1 Natural Images
Experimental setup We evaluate our algorithm using 100 images from the CelebA validation set with a resolution 256x256, normalizing all images the [0, range for quantitative analysis. All experiments incorporate Gaussian measurement noise with σg 0.01. We address the following linear inverse problems: (1) 4x super-resolution using bicubic downsampling, (2) inpainting with random mask covering 70% of missing values, (3) Gaussian deblurring with a 61 x61 kernel and a standard deviation of 3.0, and (4) box inpainting with a centered 128x 128 mask.
We present the quantitative and qualitative results of all the methods in Tab. and Fig. 3, respectively, In Tab. our method surpasses other baselines across all tasks. For more challenging tasks such as Gaussian deblurring and box inpainting, our method significantly outperforms others in terms of SSIM. Based on the MAP framework, as shown Fig. 3. our method prefers more faithful and artifact-free reconstructions, whereas others trade off for perceptual quality. We note that there is an unavoidable tradeoff between perceptual quality and restoration faithfulness [3]. Overall, our method presents higher degree of refinement. The comparison between ours and ours (w/o prior) indicates the effectiveness of the local prior term in enhancing the accuracy of the reconstructions, as evidenced by the increases in both PSNR and SSIM.
4.2 Medical application
HCP T2w dataset We utilize images from the publicly available Human Connectome Project (HCP) [50] 12-weighted (T2w) images dataset for the task of compressed sensing, which contains brain images from 47 patients. The HCP dataset includes cross-sectional images of the brain taken at different levels and angles.
8

---

TSRT
170
480 ia (000)
Poan Or Mooe plac)
48h band dnggle
16"
101
18"
10"
10
Figure 5: Ablation results of step size 7 and guidance weight X. The choice of hyperparameters for our algorithm is fairly consistent across all tasks. We choose = 10 for all experiments on CelebA-HQ For X, choose 10 for Gaussian deblurring and À 10 for the other tasks.
7588 ve
PSUB kt
24
1540
0-32 12
10
Resolution
Comproszed Sensing
Figure 6: Ablation results of iteration number K on different tasks. For super-resolution and the other three tasks, K = sufficient to achieve the best performance with the optimal step size yj and guidance weight A. However, for compressed sensing, it is necessary to increase K to obtain the best performance. We hypothesize that this due to the increased complexity of the compressed sensing operator, which requires more iteration steps to ensure the correct optimization direction.
Compressed sensing We train flow-based model from scratch on 10,000 randomly sampled images, utilizing the ncsupp architecture [48] with minor adaptations for grayscale images. We employ compression rates (1/2, 1/4, 1/10), meaning m vn. The measurement operator given by subsampled Fourier matrix, whose sign patterns randomly selected. We evaluate our reconstruction algorithm's performance on 200 randomly sampled test images.
We present the quantitative and qualitative results of compressed sensing in Tab. and Fig. 4, respectively. addition flow-based methods, we include results for two classical recovery algorithms, Wavelet [11. 32] and TV [22] priors. As shown in Tab. 2, our method outperforms the classical recovery algorithms and other flow-based baselines across varying compression rates v, demonstrating our method's capability to handle challenging scenarios and the advantages of utilizing modern generative models as priors. In Fig. 4, our method produces reconstructions that are more faithful to the original images, with fewer artifacts, leading to higher accuracy and clearer details.
4.3 Ablation studies
We use the Adam optimizer [26] for our optimization steps due to its effectiveness in neural network computations. For all tasks, we utilize N 100 steps.
Step size and Guidance weight > The use of the Adam optimizer ensures that the choice of hyperparameters, particularly the step size and the guidance weight X, remains consistent across various tasks, as illustrated in Fig. 5. Specifically, a step size of f = 10-2 2 is optimal for Inpainting (random), Inpainting (box), and Super-resolution in terms of SSIM. For PSNR, Gaussian deblurring also achieves optimal performance at 7 = 10-2, Consequently, we employ dj = 10-2 2 for all tasks. Based on the results shown in the right two subfigures of Fig. 5, select 10 for Gaussian

---

deblurring and À = 10* for the other tasks. This consistency extends to the compressed sensing experiments, where we set 10 and r 10 for all experiments involving medical images.
Iteration number K We present ablation results of the iteration number K on different tasks in Fig. We focus on the behavior of K super-resolution and compressed sensing, as performs similarly to super-resolution in the other three tasks. With the optimal choice of and À in super-resolution, provides superior performance on CelebA-HQ. A decreased step
i.e. 1) - 10-2 and = 10 size, e.g., 10-3 can help performance as K increases, but it fails to exceed the performance achieved with the optimal parameters K 1. However, for compressed sensing, it is necessary to increase K to achieve the best performance. Consequently, sensing experiments. We hypothesize that the complexity of the compressed sensing operator directly determines the number of iterations required for optimal performance.
set K 10 for all compressed
5 Conclusion
In this work, we have introduced a novel iterative algorithm to incorporate flow priors to solve linear inverse problems. By addressing the computational challenges associated with the slow log- log-likelihood calculations inherent in flow matching models, our approach leverages the decomposition of the MAP objective into multiple "local MAP" objectives. This decomposition, combined with the application of Tweedie's formula, enables effective sequential optimization through gradient steps. Our method has been rigorously validated on both natural and scientific images across various linear inverse problems, including super-resolution, deblurring, inpainting, and compressed sensing. The empirical results indicate that our algorithm consistently outperforms existing techniques based on flow matching, highlighting its potential powerful tool for high-resolution image synthesis and related downstream tasks.
Limitations and Future Work
While algorithm has demonstrated promising results, there certain limitations that suggest avenues for future research. First, our theoretical framework, built on optimal transport interpolation paths, is currently limited and cannot be applied to solve the general interpolation between Gaussian and data distributions. Additionally, order to broaden the applicability of flow priors for inverse problems, it is important to generalize the algorithm currently lacks the capability to quantify the uncertainty of the generated images, an aspect crucial for many scientific applications. It would be interesting consider approaches to post-process our solutions to understand the uncertainty inherent in our reconstruction. These limitations highlight important directions for future work enhance the robustness and applicability of our method.
approach to handle nonlinear forward models. Moreover,
10

---

Acknowledgements
The work was partially supported by NSF DMS-2015577, NSF DMS-2415226, and a gift fund from Amazon. We thank anonymous reviewers for their feedback and suggestions, which helped improve the quality of the paper.
References
[1] Muhammad Asim, Max Daniels, Oscar Leong. Ali Ahmed, and Paul Hand. Invertible generative models for inverse problems: mitigating representation error and dataset bias. Proceedings of the 37th International Conference Machine Learning, 2020.
[2] Heli Ben-Hamu, Omri Puny, Itai Gat, Brian Karrer, Uriel Singer, and Yaron Lipman. D-flow: Differentiating through flows for controlled generation. Forty-first International Conference on Machine Learning, 2024.
[3] Yochai Blau and Tomer Michaeli. The perception-distoriin tradeoff. In Proceedings of the IEEE conference on computer vision and patlern recognition, pages 2228-6237, 2018.
[4] Ashish Bora, Ajil Jalal, Eric Price, and Alexandros Dimakis. Compressed sensing using generative models. International Conference on Machine Learning, 2017.
[5] Steve Brooks, Andrew Gelman, Galin Jones, and Xiao-Li Meng. Handbook of markov chain monte carlo. CRC press, 2011.
[6] Martin Burger and Felix Lucka. Maximum a posteriori estimates in linear inverse problems with log-concave priors are proper bayes estimators. Inverse Problems, 10:114444 2014.
[7] Thorsten M Buzug. Computed tomography. In Springer handbook of medical technology, pages 311-342. Springer, 2011.
[8] Yingshan Chang, Yasi Zhang, Zhiyuan Fang, Yingnian Wu, Yonatan Bisk, and Feng Gao. Skews in the phenomenon space hinder generalization in text-to-image generation. arXiv preprint arXiv:2403.16334 2024.
[9] Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnetr, editors, Advances Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
[10] Hyungjin Chung, Jeongsol Kim, Michael T Mccann, Marc L Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. In The Eleventh International Conference on Learning Representations, ICLR 2023. The International Conference on Learning Representations, 2023.
[11] David L Donoho. Compressed sensing. IEEE Transactions information theory, 52(4))1289 1306, 2006.
[12] Bradley Efron. Tweedie's formula and selection bias. Journal of the American Statistical Association, 106(496):6002-6611 2011.
[13] Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Bocsel, et al. Scaling rectified flow transform- ers for high-resolution image synthesis. arXiv preprint arXiv:2403.0322 03200, 2024.
[14] Zhenghan Fang, Sam Buchanan, and Jeremias Sulam. What's in a prior? learned proximal networks for inverse problems. In International Conference on Learning Representations, 2024.
[15] Berthy Feng and Katherine Bouman. Variational bayesian imaging with an efficient surrogate score-based prior. Transactions Machine Learning Research, 2024.
[16] Berthy T Feng. Jamie Smith, Michael Rubinstein, Huiwen Chang, Katherine L Bouman, and William Freeman. Score-based diffusion models as principled priors for inverse imaging. In International Conference Computer Vision (ICCV). IEEE, 2023.
11

---

[17] Ian Goodfellow, Jean Pouget-Abadic Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. Communications of the ACM, 63(1 1): 1:149-144, 2020.
[18] Will Grathwohl, Ricky T. Q. Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud. Ffjord: Free-form continuous dynamics for scalable reversible generative models. In Interna- lional Conference on Learning Representations, 2019.
[19] Tapio Helin and Martin Burger. Maximum posteriori probability estimates in infinite- dimensional bayesian inverse problems. Inverse Problems, 31(8)0085009, 2015.
[20] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 23:6840-6851, 2020.
[21] Michael F Hutchinson. A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Communications in SkutisticsSSmullat and Computation, 18(3): 1059- 1076, 1989.
[22] Leonid Ruddin, Stanley Osher, and Emad Fatemi. Nonlinear total variation based noise removal algorithms. Physica D: Nonlinear Phenomena, 60:259-268, 1992.
[23] Peter A Jansson. Deconvolution of images and spectra. Courier Corporation, 2014.
[24] Tero Karras, Timo Aila, Samuli Laine, and Jaakko Lehtinen. Progressive growing of gans for improved quality, stability, and variation. arXiv preprint arXiv:1710.10190 2017.
[25] Diederik Kingma and Ruiqi Gao. Understanding diffusion objectives as the elbo with simple data augmentation. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems, volume 36, pages 65484 65516. Curran Associates, Inc., 2023.
[26] Diederík P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.69880 2014.
[27] Diederik P Kingma and Max Welling, Auto-encoding variational bayes. arXiv preprint 1312.66 14, 2013.
[28] Yaron Lipman, Ricky TQ Chen, Heli Ben- Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Eleventh International Conference on Learning Representations, 2022.
[29] Xingchao Liu, Chengyue Gong, et al. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Representations, 2022.
[30] Xingchao Liu, Xiwen Zhang, Jianzhu Ma, Jian Peng, and Qiang Liu. Instaflow: One step is enough for high-quality diffusion-bassed text-to-image generation. In International Conference on Learning Representations, 2024.
[31] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of the IEEE international conference on computer vision, pages 9730-3738, 2015.
[32] Michael Lustig, David L. Donoho, Juan M. Santos, and John M. Pauly. Compressed sensing mri IEEE Signal Processing Magazine, 25(2):72-82, 2008
[33] Morteza Mardani, Jiaming Song, Jan Kautz, and Arash Vahdat. A variational perspective on solving inverse problems with diffusion models. In The Twelfth International Conference on Learning Representations, 2024.
[34] Sachit Menon, Alex Damian, McCourt Hu, Nikhil Ravi, and Cynthia Rudin. Pulse: Self- supervised photo upsampling via latent space exploration of generative models. The IEEE Conference Computer Vision and Pattern Recognition (CVPR), June 2020.
[35] Guust Nolet. A breviary of seismic tomography. A breviary of seismic tomography, 2008
12

---

[36] Gregory Ongie, Ajil Jalal, Christopher A. Metzler, Richard G. Baraniuk, Alexandros G. Dimakis, and Rebecca Willett. Deep learning techniques for inverse problems in imaging. IEEE Journal on Selected Areas in Information Theory, 111):99-56, 2020.
[37] Ashwini Pokle, Matthew Muckley, Ricky TQ Chen, and Brian Karrer. Training-free linear image inversion via flows. arXiv preprint arXiv:23 2023.
[38] Saiprasad Ravishankar, Jong Chul Ye, and Jeffrey A Fessler. Image reconstruction: From sparsity to data- adaptive methods and machine learning. Proceedings of the IEEE, 108(():86- 109, 2019.
[39] Nicholas Rawlinson, Andreas Fichtner, Malcolm Sambridge, and Mallory Young. Seismic tomography and the assessment of uncertainty. Advances in geophysicx, 55:1-76, 2014.
[40] Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International conference on machine learning, pages 1530-1538. PMLR, 2015.
[41] François Roddier. Interferometric imaging in optical astronomy. Physics Reports, 170(2)997-
166, 1988.
[42] Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. Medical image computing and computer-assisted intervention-MICCA. 2015: 18th international conference, Munich, Germany, October 5-9, 2015, proceedings. part III 18, pages 241-241. Springer, 2015.
[43] Litu Rout, Negin Raoof, Giannis Daras, Constantine Caramanis, Alex Dimakis, and Sanjay Shakkottai. Solving linear inverse problems provably via posterior sampling with latent diffusion models. Advances in Neural Information Processing Systems, 36, 2024.
[44] John Skilling. The eigenvalues of mega-dimensional matrices. Maximum Entropy and Bayesian Methods: Cambridge, England, 1988, pages 455-466, 1989.
[45] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamics. In International conference on machine learning, pages 2256-226622 2265.22255.22 PMLR, 2015.
[46] Jiaming Song, Arash Vahdat, Morteza Mardani, and Jan Kautz. seeudoinverse-guidee diffusion models for inverse problems. In International Conference Learning Representations, 2023.
[47] Yang Song, Conor Durkan, lain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models, Advances in neural information processing systems, 34:1415- 1428, 2021.
[48] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13455 2020.
[49] Paul Suetens. Fundamentals of medical imaging. Cambridge university press, 2017.
[50] David C Van Essen, Stephen M Smith, Deanna M Barch, Timothy EJ Behrens, Essa Yacoub, Kamil Ugurbil, Wu-Minn HCP Consortium, et al. The wu-minn human connectome project: an overview. Neuroimage, 80:62-79, 2013.
[51] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
[52] Marinus T Vllaaddiigerrrrkkk and Jacques A Boer. Magnetic resonance imaging: theory and practice. Springer Science & Business Media, 2013.
[53] Zhou Wang, A.C. Bovik, H.R. Sheikh, and E.P. Simoncelli. Image quality assessment: from error visibility to structural similarity. IEEE Transactions on Image Processing, 13(4):600-61,, 2004.
13

---

[54] Jay Whang, Erik M. Lindgren, and Alexandros G. Dimakis. Composing normalizing flows for inverse problems. Proceedings of the 38th International Conference on Machine Learning. 2021.
[55] Jianwen Xie, Yang Lu, Song-Chun Zhu, and Yingnian Wu. A theory of generative convnet. In Maria Florina Balcan and Kilian Q. Weinberger, editors, Proceedings of The 33rd International Conference Machine Learning, volume 48 of Proceedings of Machine Learning Research, New York, New York, USA, 20-22 Jun 2016. PMLR.
pages
[56] Hanshu Yan, Xingchao Liu, Jiachun Pan, Jun Hao Liew, Qiang Liu, and Jiashi Feng. Perflow: Piecewise rectified flow as universal plug-and-play accelerator. arXiv preprinl arXiv:2405.075100 2024.
[57] Peiyu Yu, Dinghuai Zhang, Hengzhi He, Xiaojian Ma, Ruiyao Miao, Yifan Lu, Yasi Zhang, De- qian Kong, Ruiqi Gao, Jianwen Xie, al. Latent energy- based odyssey: Black-box optimization via expanded exploration in the energy-based latent space. arXiv preprint arXiv:2405.167708 2024.
[58] Yasi Zhang, Peiyu Yu, and Ying Nian Wu. Object-conditioned energy-based model for attention map alignment in text-to-image diffusion models. In Synthetic Data for Computer Vision Workshop @ CVPR 2024, 2024.
[59] Yasi Zhang, Peiyu Yu, and Ying Nian Wu. Object-conditioned energy-based attention map alignment in text-to-image diffusion models. In Aleš Leonardis, Elisa Ricci, Stefan Roth, Olga Russakovsky, Torsten Sattler, and Güil Varol, editors, Computer Vision ECCV 2024, pages 55-71, Cham, 2025. Springer Nature Switzerland.
[60] Yaxuan Zhu, Zehao Dou, Haoxin Zheng, Yasi Zhang, Ying Nian Wu, and Ruiqi Gao. Think twice before you act: Improving inverse problem solving with meme. arXiv preprint arXiv:2409.0m 08551, 2024.
14

---

Appendix
A Proof
Before we dive into the proof, we provide the following three lemmas.
Lemma 1. Consider a vector-valued function f : [0. 1] + RR. Then for any t E [0, 1}. we have that
11, f())ds
2 - L [1(())))ds.
Proof. For each E [0, let fi(s) E I. denote the i-th component of f(s). Recall Jensen's
inequality: for any convex function R - R and integrable function h:: [0, + R. we have
(
[." g(h(t)dd..
Using convexity of the function
+ t and applying Jensen's 's inequality, we see that
1.- M(s)ds
Σ So. M()dds)
L: J((s)
fs(s)
L. [1]())])
Lemma 2 (Tweedie's Formula [12]). 1f ~ g()), =1u ~ M(aji,021) and therefore ~ s(-), we have
[[p|2]= = - + o² - log f(=)]-
(12)
Lemma Suppose y A(x,) + E where x. I (20) with xo being the solution to Eq. (9), A: R + IR is linear, ~ (0, o and Ia exactly follows the path It Og + 3,00 for any
time E 10. Then we have
p(yolx)) (Nx,,, 1I,
(13)
and hence
log p(yla(m))) log ((g(()))):)))))))))))))))))))))) + log(o?), Vl.
(14)
Proof. Recall that the generated auxiliary path yt = a.y + B,Azo. By assumption, we have A(x;) A(axx + B,z0) A(x(x))) + Aro. By subtracting these two equations, we have V1- A(x,) = an - A(x(xo)).
(15)
As y)())))) ~ (Ar, I)) we have 3e jz, ~ (Arr, ata?l) The proof for Eq. (13) is done. Next, we examine the log probability as follows:
11 Ar log log(2Ta(o
p(v((())
(16)
20 3
||a -((x(10))12
m
log (2
(17)
2070z
2
lly A(x(2))l1
log(2To7)
log(a?)
(18)
20
2
= logr(ylo(20))
2 log(a?)
(19)
15

---

A.1 Proof of Proposition 1
Trained by the objective defined in Eq. (5), the optimal velocity field would be
(0 t) E(dgz1 (,
(20)
E(áxT + 3,
tt
Br x())
axæ
Given Ie. 20
2t
8r.
(21)
3c
(0: - Be
Ott
E( (22)
+
(22)
Bt
3t
(0z Br
Ot
1
(Iz + 3i
log (gg((())))
Be dx. # Lemma 2(Tweedie's Formula) B
Be
(23)
By defining the signal-to-noise ratio as A, ar/S, and rearranging the equation above, get exactly Eq. (11) which we display again below:
(dlog de
d dloog Br 21) (o((ss,))
V log Bi
(24)
dt
dt
When 04 1, the equation above becomes
1
Vx, log p(xe)
= - (-2, (x)(x)))))
(25)
A.2 Proof of Theorem
Before we dive into the proof, we first point out lim Σ === % = 1. Define the timestep
t At. Conversely, 1+ t/ At is function of In this sense, we define the i-th step
Riemannian discretization of the integral So ir (va((ss)) dl Ape
(Do((e,)) 8x
Al.
We first decompose the global MAP objective as follows: S."
logp(z(20)lv) log
dt logp()lz(20)) log
(26)
Ox
N
N
lim At
%s logg(po) + lim A4-10
Api
(27)
+ lim
74 Poogp(Via + c.l - p())
(28)
At
where the decomposition of the second term utilizes the property of the discretization of Riemann integral, and that of the third term utilizes the result in Lemma 3 and thus
7 (o))))))). By the property of limits, i.e. o T Ap:)
limar-o(E)) %) limar-oo (E Api) = limat-o Σi Api, we can further decompose the second term in Eq. (28) into lim (E> ( Ap.).
By extracting the limit out in Eq. (28), the equation becomes
lim Ar- {1) (logp(xo) + Ap1 + log + ta
+ 72 [log
Ap1 + Ap2 + log logp(vasr + 02)
+
+ YN [logp(zo) Api + Ap2 + ... ApN + logp(wNA)/Zwax) eN]
+ Ap2 + (71 + 72 72)4ps p3 +... + (n1 + 72 + + x Apnl logp(z))
(29)
N MrJ:
lim
71
Ap, + Σ
74 logp(y)
(30)
1=1
1=
1=1
16

---

where J. == logp(ro))+ E=== Ap, + tog We further define the c(N) = E 7:c. logp(y) p(y).
Recall that J, = log (( - tr inequality, have
At + log By triangle
8x
N nJ: -
log logp(z(zo)ly)
(31)
N Σ
N
N
(ogp(x((oolv))
nh: - c(N)
+
Σ
M
Σ
(32)
Taking the limit on both sides, we have
N Σ
lim A1-0
log (gg((((F))l))
mJ. c(N)
(33)
1=1
N
lim log logp(<(zo)lv))
Σ
nJ c(N)
+ lim
7:J
(34)
1=1
In the following, we analyze the two terms on the right-hand side one by one. For the first term: as 1 :: IR a continuous function, the first term on the right-hand side is equal to
log logp(x(<ro)1)
lim A1-0
MD. c(N)
(35)
N
lim
Apj
(36)
A1-0
1=2
1==
N
1
lim
Apj
(37)
At-+0
-N
2N
j=2
Ap;
App
(38)
A1-0
2N 2N-
A1-0
2
j=2
j=2
where the first equation derived by subtracting the first term in Eq. (30) from Eq. (26). As the 3
velocity field Vo R" R R² satisfies sup.ER", sejo.) |tr
on vo(2, s) I C₁ for some universal
constant C1, we have AApl S At. The first term in (38) would be
1
1
N -++1
AP
SCA
Σ j=2
2N-J+1
SCA O(At).
(39)
Similarly, the second term in (38) would be
N
NN11
App
CAA= Ci
--= O(A))
(40)
2
2
2N
J=2
Combining the results in Eq. (39) and Eq. (40), we get
logp(æ(20)1o) - lim Σ -
c(N) =0.
(41)
At
17

---

For the second term: Intuitively, the error between the integral and the Riemannian discretization goes 0 as At tends to O. Rigorously,
N
N
lim At-0
Mi
lim Ar-0
Ml..
Fi)
(42)
2m
N
s-ar
(va (1g(xs)))
1-1
lim At-0
Σ
%
tr
ds
Apj
(43)
dr
iml
jm1
N
2-1
lim
%
tr
ds
Σ
Apj
0.
(44)
A0-0
Ox
j=1
Combining the results of the first term and the second term, we get the proof of theorem 1 done.
B Compliance of Trajectory
To quantify our deviation from the assumption of having z exactly follow the interpolation path a, + 3, we define the following: given a differentiable process {ze} and an interpolation path specified by a == {ar} and B == {Be). we define the trajectory's compliance Sa,p((z))) to the interpolation path as
Sa,z(1z))) == L.'' Ep(co))((xx) || - (á + 312 dt.
(45)
This generalizes the definition of straightness in [29] to general interpolation paths. We recover their definition by setting a = and 3, = 1 t. In certain cases, we have exact compliance with the predefined interpolation path. For example, when {2+} is generated by up and ax and 8, = 1- t, note that Sa ( = is equivalent to (= f) = c where c is a constant, almost everywhere. This ensures that 21 2u In this case, when generating the trajectory through ODE solver with starting point zy and endpoint 24, we have I2 = agz + S, Vt. When Sa to O, show in Proposition 2 that can bound the deviation of our trajectory from the interpolation path using this compliance measure. When specifying our result to Rectified Flow, we can obtain an additional bound showing that when using I--Reciffidd Flow, the deviation of the learned trajectory from the straight trajectory is bounded by (l)
(2+}) is not equal
Proposition 2. Consider a differentiable interpolation path specified by a == (ar) and B = ( Then the expected distance between the learned trajectory 2 =o + So vo volzs, s)ds and the predefined trajectory 2 = 20 + So (à,z2 + B,20)ds can be bounded as
(p( [1l - 2x S Sa,s((z))).
(46)
If the differentiable process {zt} is specified by I.-Rectified Flow and ar = and Br = 1 t for all
t [0, 1], then we additionally have
Eptz))p(()) [1l2, - 2el33 < (
(47)
Proof. At time f, we are interested in the distance between real trajectory 2, = 2g + So vo 0o s)ds and preferred trajectory & zo + 1n (à (322))s.. Using the result in Lemma 1, the distance can be bounded by
|| - =el12 = 1,,- (u0(20)8) -(0sz1 - 2
2
(48)
< f
[100(ze1 8) - (á BB20))|dds.
(49)
18

---

Figure 7: Ablation results of K in terms of SSIM on different tasks.
Therefore,
Bp( 2el2 s (p(2o)p((s) [ [1oo(2).)) (á,21 -(d,)) 21
(50)
Bptco)stey)lloozsss)) (ávzz) - B1z0)]1dds
(51)
Bptew) (l - (dr.21 B.20)
(52)
= (az)((z))
(53)
If (2,, E [0, 1}} is learned T-reetfied flow, i.e. 04 and B = in this case, where the times of rectifying the flow, by Theorem 3.7 in [29], we have Sa,a((=)) = O(1/T) and thus
pp(=)))((1)) & 21 O(1LL)
(54)
Empirically, [30, 29] found = 2 generates nearly straight trajectories for high-quality one-step generation. Hence, while this result gives simple upper bound, in practice the trajectories may comply more faithfully with the predefined interpolation path than this result suggests.
C
Additional Results
C.1 Additional Ablations
Iteration steps
We provide additional ablation results of K in terms of SSIM in Fig. 7.
NFEs We first refer to Fig. 2(c) for a preliminary ablation on N using a toy example. Next, we show PSNR and SSIM scores for varying N in the task of super-resolution We find that N 100 is the best trade-off between time and performance. The ablation results are shown in Fig. 8.
265
Figure 8: Ablation results of the NFEs N on the super-resolution task.
D Computational Efficiency
In Tab. 3, we present the computational efficiency comparison results. Note that OT-ODE the slowest as it requires taking the inverse of a matrix rÇAAT + o each update time. Our method requires taking the gradient over an estimated trace of the Jacobian matrix, which slows the computation.
19

---

Table 3: Computational time comparison. We compare the time required to recover 100 images for the super-resolution task on single GPU.
DPS-ODE
OT-ODE
Ours (w/o prior)
Ours
Time(h) 0.36
4.10
0.83
2.72
E Implementation Details
Experiments were conducted on Linux-based system with CUDA 12.2 equipped with 4 Nvidia R9000 GPUs, each of them has 48GB memory.
Operators For all the experiments on the CelebA-HQ dataset, we use the operators from [10]. For all the experiments on compressed sensing, we use the operator CompressedSennsnnn@ppraaaa defined in the official repository of [14]
Evaluation Metries are implemented with different Python packages. PSNR is calculated using basic PyTorch operations, and SSIM is computed using the pytorch msssim package.
E.1
Toy example
The workflow begins with using 1,000 FFHQ images resolution of 0224x1024 These images are then downscaled to 16x 16 using bicubic resizing. A Gaussian Mixture model is applied to the downsampled images, resulting in mean and covariance parameters. The mean values transformed from the original range of [0,1] to [-1.1]. Subsequently. 10,000 samples are generated from this distribution to facilitate training score-based model resembling the architecture of CIFARIO DDPM++ The training process involves 10,000 iterations, each with batch size of 64, and utilizes the Adam optimizer [26] with learning rate of 2c-4 and warmup phase lasting 100 steps. Notably, convergence is achieved within approximately 200 steps. Lastly, the estimated log- log-likeldhood computation for a batch size of 128 takes around minutes and 30 seconds. We show uncured samples generated from the trained models in Fig. 9.
Figure 9: Generated samples from the flow trained on 10,000 Gaussian samples.
E.2 Medical Application
In this setting, o, 0.001. We the nesspp architecture, training from scratch 10k images for 100k iterations with a batch size of 50. We set the learning rate to 1 x 10-2, Sudden convergence appeared during our training process. We use 2000 warmup steps. Uncured generated images are presented in Fig. 10.
'nttps://githan www.aauu...mmmmaeememeeeeeeeeeee com/Sulam Group/learned- hooupp/oarnee-proximll---ee proximal- wat.aaaalennttoorreeeeeeeeaaaaa
20

---

Figure 10: Generated samples from the flow trained on 10,000 HCP T2w images.
E.3 Implementation Baselines
OT-ODE As OT-ODE [37] has not released their code and pretrained checkpoints. We reproduce their method with the same architecture as in [29]. We follow their setting and find initialization time (' has great impact on the performance. We use the y init method in their paper. Specifically, the starting point
Tv =ty (1 t'),, NO,
(55)
where l' is the init time. Note that in the super-resolution task we upscale y with bicubic first. We follow the guidance in the paper and show the ablation results in Fig. 11 and Fig. 12. Inpainting(random)
Resolutin
Gaussian Deblurring
Inpainting(box)
Figure Hyperparameter 1' selection results for OT-ODE on the CelebA-HQ dataset. We select 1' 0.2, 0.1, 0.2, 0.2 for super-resolution, inpainting(random) Gaussian deblurring, and inpainting(box). respectively.
DPS-ODE
We use the following formula to update for each step in the flow:
u(r2, y) v(x,) G. (- (-Vz, -aly--Ai.. A 12).
where G. is the step size to tune. We refer to DPS for the method to choose G. We set G. 21y- ( We demonstrate the ablation of n for this baseline in Fig. 13 and Fig. 14. Note that there significant divergence in PSNR and SSIM for the task of inpainting (box). As we observe that artifacts are likely appear when 2 100, we choose the optimal 75 for the best tradeoff.
21

---

2-0
Figure 12: Hyperparameter 1' selection results for OT-ODE on the HCP T2w dataset. We select
t' = 0.1 for all the experiments.
RED-Diff and IIGMM We use the official repository from Nvidia to reproduce the results of RED-Diff and IIGDM with the pretrained CelebAHQ checkpoint using the architecture of the guided diffusion repository" from OpenAI.
For RED-Diff, the optimization objective min, |y - A(p) + ((sg(en(x,,)) - +)) p.. Fol- lowing the implementation of the original paper, tasks. We choose learning rate ir = 0.25, À = 0.25 for super-resolution, inpainting(random) and inpainting(box) and Ir 0.5, À 0.25 for deblurring as recommended by the paper.
Adam optimizer with 1,000 steps for all
For IIGDM, we follow the original paper and use 100 diffusion steps. Specifically, we use rj = 1.0
oi are used if there is an improvement which corresponds to the VE-SDE. Adaptive weights r = Ito; on metrics.
Wavelet and TV priors We use the pytorch package DeepInverse' to implement Wavelet and TV priors. For both priors, we use the default Proximal Gradient Descent (PGD) algorithm and perform grid search for regularization weight À the set (10", 10-1 10-2 10-3 10 +] and gradient stepsize
in {10', 10 10 10 10 10 10=4). The maximum number of iteration is 3k, 5k, and 10k for compression rate v 1/2,1/4 and 1/10, respectively.
stopping criterion is the residual norm 1a all < 1 x 10-5 and the initialization of the algorithm is the backprojected reconstruction, ie., the pseudoinverse of AA applied to the measurement y.
|
For the TV prior, the objective we aim minimize is ming 11 - y112 \| We find that the optimal combination of hyperparameters is À 0.01, n) = 0.1 for all the values of v.
For the Wavelet prior, the objective want minimize is min, | v113 A| Marl/2 We the default level of the wavelet transform and select the "db8" Wavelet. The optimal combination of hyperparameters is
0.1, n 0.1 for all the values of
Pattpe://githuu..com// http://www...ooeaaa//a/ddddduuuu
22

---

Resolution
Inpainting(random)
Gaussian Deblurring
Inpainting(box)
1100
13.00
410
12
135
133
18
0110
3750
Figure 13: Hyperparameter 7) selection results for DPS-ODE. We select 1) = 1000, 750, 200, 75 for super-resolution, inpaintingrrandom) Gaussian deblurring, and inpainting(box), respectively.
--
Figure 14: Hyperparameter selection results for DPS-ODE on the HCP T2w dataset. We select " 200 for all the experiments.
NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
The claims made in the abstract match theoretical and experimental results.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]
23

---

Justification:
They are discussed in Section Limitations.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model vell-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification:
They are provided in Section Method.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: They are provided in Section Experiments and Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.

---

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g.. with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibilityy In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We have provided sufficient implementation details and links for original repositories.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: /ttpsscccppbblcc////eeccppcpppcp wwwubblccguiideee guwdee/ /CodeSubmii ssionPolicy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
25

---

Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification:
They are provided in Section Experiments and Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: Standard deviations are provided in Tables.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by enror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
They are provided in Appendix.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
26

---

The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethies
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: http://eerippppeeeee cwtprriss.ccppppppocm cc/public/ ethicshuudelines
Answer:
[Yes]
Justification:
We confirm that the paper conforms with NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g.. if there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: They discussion in Section Broader Impacts.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Answer: [NA]
Justification: The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses such risks.
27

---

Released models that have high risk for misuse or dual use should be released with necessary safeguards to allow for controlled use of the model. for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: They are properly cited.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
They are well documented in Section Experiments and Appendix.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects.
28

---

Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
29


Reviews:
[{"review_id": "b6YXjgkAvc", "reviewer": "Reviewer_hhXk", "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.", "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?\n*    In terms of runtime, how does your method fair with other types of MAP estimates employing different generative models as priors? \n*    What changes will it take to apply your method to non-linear inverse problems?", "limitations": "The authors were upfront about the limitations of their method and stated these in Appendix A.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "*    The paper proposes a theoretically motivated approximation of MAP estimates utilizing flow-matching priors.\n*    The paper is well written and overall well structured.\n*    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\n*    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. \n*    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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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:\n\n**The use of log-likelihood and flow prior** \n\n**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. \n\n**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.\n\n**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.\n\n**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. \n\n**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. \n\n 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. \n\n[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)\n\n[2] Training-free linear image inverses via flows. Ashwini Pokle, Matthew J. Muckley, Ricky T. Q. Chen, Brian Karrer. TMLR 2024. \n\n[3] A Variational Perspective on Solving Inverse Problems with Diffusion Models. Mardani, Morteza and Song, Jiaming and Kautz, Jan and Vahdat, Arash. ICLR 2024.\n\n[4] Pseudoinverse-guided diffusion models for inverse problems. Song, Jiaming and Vahdat, Arash and Mardani, Morteza and Kautz, Jan. ICLR 2023.\n\n[5] Compressed Sensing using Generative Models. Ashish Bora, Ajil Jalal, Eric Price, and Alexandros G. Dimakis. ICML 2017.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "sKUCYpjjB8", "reviewer": "Reviewer_Xm5z", "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.\n\nTo 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.\n\nThe 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.\n\nThe 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.", "questions": "NA", "limitations": "NA", "rating": 7, "confidence": 1, "soundness": 3, "presentation": 3, "contribution": 3, "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\n- 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.\n- clarity: fairly clear, improvements on the flow of argument would be desired.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PfOWxZuKDR", "reviewer": "Reviewer_sT2G", "summary": "The paper proposes a flow prior under the MAP structure for solving inverse problems, a theoretical analysis is also given.", "questions": "N/A", "limitations": "N/A", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n2. Incorporating the flow prior into MAP to solve inverse problems is interesting and worth exploring.\n3. 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$.\n\n2. 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.\n\n3. 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. \n\n[1] Donoho, David L. \"Compressed sensing.\" IEEE Transactions on information theory 52.4 (2006): 1289-1306.\n\n[2] Lustig, Michael, et al. \"Compressed sensing MRI.\" IEEE signal processing magazine 25.2 (2008): 72-82.\n\n4. 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.\n\n[1] Mardani, Morteza, et al. \"A variational perspective on solving inverse problems with diffusion models.\" arXiv preprint arXiv:2305.04391 (2023).\n\n5. 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.\n\n6. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank you for your helpful and detailed feedback, which significantly helped us improve the paper quality. \n\n**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.\n\nThe $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. \n \n\n**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. \nImplementation details are included at the end of the [Author Rebuttal] above. \n\n**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.  \n\n**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.\n\n**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).\nThe 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.\n\nWe 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.\n\n\n[1] A Variational Perspective on Solving Inverse Problems with Diffusion Models. Mardani, Morteza and Song, Jiaming and Kautz, Jan and Vahdat, Arash. ICLR 2024. \n\n[2] Pseudoinverse-guided diffusion models for inverse problems. Song, Jiaming and Vahdat, Arash and Mardani, Morteza and Kautz, Jan. ICLR 2023. \n\n[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)", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "F9aYPE5Y9O", "reviewer": "Reviewer_qjx2", "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.", "questions": "- Does the proposed method work for very noisy corruptions with high measurement noise (e.g., sigma_y = 0.2 or more)?\n- Does the proposed method work for severely ill-posed inverse problems like 8x super-resolution or inpainting with 90% of pixels missing?\n- Why was the algorithm evaluated on only 100 images from the CelebA-HQ dataset? Why wasn't the entire validation dataset used?\n-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?", "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.\n- The presentation of quantitative results could be improved by using tables for easier comparison and interpretation.\n- The evaluation metrics need to be expanded to include FID and LPIPS scores to provide a more comprehensive assessment of performance.\n- Blurry results", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The paper introduces a unique method for solving inverse problems and reconstructing a single image.\n- The proposed method is supported by a solid theoretical foundation.\n- 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.\n- Need for Additional Metrics: Metrics like FID and LPIPS scores should be included alongside PSNR and SSIM.\n- Figure 2 Presentation: The data in Figure 2 should be presented in a table format for easier interpretation of quantitative performance improvements.\n- 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.\n- 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.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for your helpful review and provide the additional experiments which we hope address your concerns:\n\n**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. \n\nThank 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. \n\n**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.\n\n**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.\n\n**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.\n\n**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.   \n\n\nWe 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.\nWe kindly request a re-evaluation of the rating if some of your concerns have been resolved.\n \n\n\n \n\n\n[1] A Variational Perspective on Solving Inverse Problems with Diffusion Models. Mardani, Morteza and Song, Jiaming and Kautz, Jan and Vahdat, Arash. ICLR 2024. \n\n[2] Pseudoinverse-guided diffusion models for inverse problems. Song, Jiaming and Vahdat, Arash and Mardani, Morteza and Kautz, Jan. ICLR 2023. \n\n[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)", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
Dear Reviewers, 

We first thank you for your valuable feedback and appreciate the recognition of our paper's strengths, such as 

- concise presentation from concepts to methods and clear motivation (sT2G),
- a solid theoretical foundation supporting the method (qjx2),
- consistently improved performance in distortion across multiple tasks and datasets (hhXk).

In the attached pdf, we provide additional experimental results that we hope address your concerns (implementation details are documented at the end):


**[Highlight: more diffusion-based baselines in Table 2]** 
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.



**Table 1(two classical recovery algorithms and a more challenging setting $\nu=1/10$ included)**: Results of compressed sensing with varying compression rate $\nu$ on the HCP T2w dataset. We have added results for two classical recovery algorithms, Wavelet and TV priors, as well as a more challenging setting  $\nu=1/10$, where only 10\% of the output signal is observed. 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.

  


**Table 2(two new baselines and a highly noisy setting $\sigma_y=0.2$ included)**: Quantitative comparison results of PSNR and SSIM on the CelebA-HQ dataset (best values highlighted in blue and second-best underlined). We have included two representative diffusion-based baselines: 1)RED-Diff [1], a variational Bayes-based method; 2)$\Pi$GDM [2], an advanced MCMC-based method. Our method is competitive with diffusion-based baselines and other approaches in terms of recovery fidelity. Additionally, we included an extremely noisy setting with 
 $\sigma_y=0.2$.  The last row of the table shows our method's capability to handle noisy settings across different tasks. 


**Table 3(FID and LPIPS scores included)**: Quantitative comparison results of FID and LPIPS on the CelebA-HQ dataset (best values highlighted in blue and second-best underlined). This table shares the same setting as Table 2. Our method is also competitive with other baselines in terms of perceptual quality as evidenced by LPIPS.

Please see our reviewer-specific feedback for more information.  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)

---
[Implementation Details]

Table 1: We use the pytorch package *DeepInverse*$^1$ to implement Wavelet and TV priors as shown in Tab.1 of the pdf. For both priors, we use the default Proximal Gradient Descent  algorithm and perform a grid search for regularization weight $\lambda$ in the set {$10^0, 10^{-1}, 10^{-2}, 10^{-3}, 10^{-4}$} and gradient stepsize $\eta$ in {$10^0, 10^{-1}, 10^{-2}, 10^{-3}, 10^{-4}$}. The maximum number of iterations is 3k, 5k, and 10k for compression rate $\nu = 1/2, 1/4,$ and $1/10$, respectively. The stopping criterion is the residual norm $\frac{||x_{t-1}-x_t||}{||x_{t-1}||} \le 1\times 10^{-5}$ and the initialization of the algorithm is the backprojected reconstruction, i.e., the pseudoinverse of $\mathcal{A}$ applied to the measurement $y$. 

- For the TV prior, the objective we aim to minimize is $\min_x \frac{1}{2}||\mathcal{A}x - y||_2^2 + \lambda \||x\||{\rm tv}$.  We find that the optimal combination of hyperparameters is $\lambda = 0.01, \eta = 0.1$ for all values of $\nu$.     

- For the Wavelet prior, the objective we want to minimize is $\min_x \frac{1}{2}||\mathcal{A}x - y||_2^2 + \lambda ||\Psi x||_1$. We use the default level of the wavelet transform and select the “db8” Wavelet. The optimal combination of hyperparameters is 
$\lambda=0.1, \eta=0.1$ for all values of $\nu$.

Table 2: We use the official repository$^2$ from Nvidia to reproduce the results of RED-Diff and $\Pi$GDM with the pretrained CelebAHQ checkpoint using the architecture of the guided diffusion repo from OpenAI$^3$. 

- For RED-Diff, the optimization objective is $\min_\mu  ||y - \mathcal A (\mu)||^2 + \lambda  (sg(\epsilon_\theta(x_t,t) - \epsilon))^T \mu$. Following the implementation of the original paper, we use Adam optimizer with 1,000 steps for all tasks. We choose learning rate $lr=0.25, \lambda=0.25$ for super-resolution, inpainting(random) and inpainting(box) and $lr=0.5, \lambda=0.25$ for deblurring as recommended by the paper. 

- For $\Pi$GDM, we follow the original paper and use 100 diffusion steps. Specifically, we use $\eta = 1.0$ which corresponds to the VE-SDE. Adaptive weights $r_t^2 = \frac{\sigma_{1-t}^2}{1+\sigma_t^2}$ are used if there is an improvement on metrics.  


1. https://deepinv.github.io/deepinv/
2. https://github.com/NVlabs/RED-diff
3. https://github.com/openai/guided-diffusion

Meta-review: The paper has received all positive reviews, with one strongly in favour but with low confidence and the rest borderline accepts. After the discussion phase, reviewers raised their score. Reviewers appreciate the theoretical contribution of integrating flow matching with into MAP for solving inverse problems is a well motivated and interesting approach, the good writeup and improved experimental results. On the other hand, some significant weaknesses also persist, including lack of some important metrics (perceptual quality), limited usefulness of likelihood information, limited number of baselines and experimental datasets and asymptotic error bound usefulness
In light of the positive rebuttal phase that successfully addressed some of the above-mentioned issues and the positive scores and the overall solid contribution, the paper is accepted.
Decision: Accept (poster)