Flow Matching for Posterior Inference with Simulator Feedback

• In the strong lensing section, while a comparison with HMC is done, a comparison with "traditional" neural simulation-based inference approaches like NPE is lacking. This comparison would significantly strengthen the outcomes of this experiment.
• The pretraining/finetuning tradeoff is not comprehensively explored -- as far as I understand, a big advantage of the method is that one could just finetune using a smaller number of simulation calls. While this is mentioned briefly and tested for specific cases, a more comprehensive study of how fraction of finetuning for a fixed simulator budget affects the outcome would make the results quite a bit stronger.
• The high-level presentation could be made slightly clearer -- e.g., in Fig. 1, marking that the it is the posterior $p(\theta\mid x)$ that is being modeled/targeted, to immediately situate the reader to the problem setting.

• The empirical evaluation seems quite a bit limited. Specifically, the paper only considers a single applicaiton on strong gravitational lensing. The other synthetic tasks are quite small and it is unclear how general the method is. Moreover, the results on the gravitational lensing experiment, the results are not convicing. The residuals indicate the samples are not capturing the posterior. Finally, the authors consider a relatively simple variant of the problem (23D). In this setting FM + simulator feedback just acts as a faster approximation to MCMC. Where FM + Simulator feedback might have an advantage is high dimensional unstructured data (e.g. images in the case of lensing)
• Another shortcoming of the empirical evaluation is the relatively limited baselines. There are several approaches for unbiased inference with diffusion priors (like DPS) [1-4], so it would be good to add comparisons to some of these baselines.
• There is some missing discussion about related work about the guidance of flow matching models [5-6].
• The code to reproduce experiments is not provided though there are sufficient details in the paper.

[1] Wu, Z., Sun, Y., Chen, Y., Zhang, B., Yue, Y., & Bouman, K. L. (2024). Principled Probabilistic Imaging using Diffusion Models as Plug-and-Play Priors. arXiv preprint arXiv:2405.18782.

[2] Wu, L., Trippe, B., Naesseth, C., Blei, D., & Cunningham, J. P. (2023). Practical and asymptotically exact conditional sampling in diffusion models. Advances in Neural Information Processing Systems, 36.

[3] Dou, Z., & Song, Y. (2024). Diffusion posterior sampling for linear inverse problem solving: A filtering perspective. In The Twelfth International Conference on Learning Representations.

[4] Chung, H., Lee, S., & Ye, J. C. (2023). Decomposed Diffusion Sampler for Accelerating Large-Scale Inverse Problems. arXiv preprint arXiv:2303.05754.

[5] Nisonoff, H., Xiong, J., Allenspach, S., & Listgarten, J. (2024). Unlocking Guidance for Discrete State-Space Diffusion and Flow Models. arXiv preprint arXiv:2406.01572.

[6] Zheng, Q., Le, M., Shaul, N., Lipman, Y., Grover, A., & Chen, R. T. (2023). Guided flows for generative modeling and decision making. arXiv preprint arXiv:2311.13443.

Despite its strengths, the paper has a few major and minor weaknesses.

Major weaknesses:

1. I’m really concerned about the presentation of the results and the abilities of this method, because of the lack of enough comparisons and metrics.
2. Regarding the gravitational lensing problem (presented as the main real-world task being tackled), they are operating in a relatively small space, making the problem easy and solvable. However, they didn’t include any coverage test (e.g., TARP [1], or any other), so it’s hard to say if the posteriors are good. Moreover, as far as I understand, the evaluation test is only on their own simulations, the one that the model were being trained on. In particular, we don’t know if the model is robust to any OOD examples. Finally, the presented residuals (e.g., in Fig. 6) are looking bad.
3. Regarding the part: „(…) however, previous methods are usually restricted to point estimates, use simple variational distributions or Bayesian Neural Networks (Schuldt et al., 2021; Legin et al., 2021; 2023; Poh et al., 2022) that are not well suited to represent more complicated high-dimensional data distributions.” - Legin et al. 2021; 2023 use a likelihood-free inference (or simulation-based inference) framework to get posteriors from simple feed-forward nn (not Bayesian).
4. The authors proposed the intensive tests only on LV, which is relatively simple problem and for sure not enough for fair comparison. Moreover, the results in Tab. 1 don’t show any superiority of the proposed method - in particular, NSF is getting similar or better results.

Minor weaknesses:

1. The authors should include comparisons with other novel posterior sampling baseline methods than DPS, e.g., LGD−MC [2].
2. In line 319, should be „spline”

References:

[1] Lemos, P., Coogan, A., Hezaveh, Y., & Perreault-Levasseur, L. (2023, July). Sampling-based accuracy testing of posterior estimators for general inference. In International Conference on Machine Learning (pp. 19256-19273). PMLR.

[2] Song, J., Zhang, Q., Yin, H., Mardani, M., Liu, M. Y., Kautz, J., ... & Vahdat, A. (2023, July). Loss-guided diffusion models for plug-and-play controllable generation. In International Conference on Machine Learning (pp. 32483-32498). PMLR.

While I must acknowledge certain positive aspects of the paper, I have also several concerns that motivate my negative recommendation, which I am listing below.

1. While I find the figures and tables quite enlightening, I find the presentation being a weakness. There are multiple hand-wavy explanations and claims that I find a bit confusing. See below for concrete examples.
2. Empirical validation: I was surprised by the numbers from figure 3 and 4 which seems worse than existing alternatives. In particular, the paper only compares to offline methods while it is clear from 1 that SNLE or SNPE which are sequential alternatives to NPE/NLE and perform much better than the proposed method. It is also arguable whether these benchmarks are the most relevant ones as there are quite low dimensional in term of observation dimensionality and are not fully representative of applications where SBI can shine. For instance, gravitational waveforms and spiking neurons are interesting benchmark used in many SBI papers. I also find the results of Flow matching + posterior quite bad compared to MCMC methods in figure 14 and 15 where we can clearly see an issue with bias and also variance of the predicted posteriors. Overall, the empirical validation seems insufficient to me and do not clearly demonstrate that the method proposed is of any real use.
3. Relevance: Sequential SBI methods, which call the simulator online, are often complicated to motivate as they require to access the simulator while also arguing doing inference on this simulator without SBI is hard because evaluating the simulator is computationally demanding. While I agree this is not the case for all applications, I would expect the paper to clearly highlight and benchmark the method on use cases where simulating more samples offline to get a good amortised posterior is not enough and calling the simulator online does not take too much time.

I am not confident these concerns can be addressed in the scope of the discussion (especially regarding results and presentation) but I am open to discussing with authors.

Hand-wavy explanations examples:

• l33-36: Seems an intricated way of explaining likelihood and posterior distributions which are very standard mathematical objects most reader should already know about.
• l37-43: The presentation of SBI is again a bit weird in my opinion. It does not clearly say when and why SBI may be necessary and what it solves. It may be interpreted as if SBI was only about Bayesian inference for the uninformed reader where frequentist methods exist as well.
• l48: what do you mean by "it became clear ... be specified a priori"?
• l49-50: this a vague and pretty strong statement that I would expect to be clarified and supported by reference.
• l52-54: I do not understand what you mean by saying there is "no feedback loop".
• l74-78: you mix the high level description of the method with specific implementation decisions which makes it hard to understand on what aspects the reader should focus on.
• Section 3 was presented in a way that I did not find easy to digest.
• l205: "a fundamental problem..." why is it a fundamental problem is unclear to me.
• l254-255: What shall we conclude from that sentence?
• l277-282: It seems like a hacky solution
• l283: While I understand you are trying to emphasise that the algorithm still learn the "true" posterior distribution, this is stated in a hand-wavy way in my opinion.
• l345-351: This is quite confusing again. Why is there a dropout layer? It is quite complicated to grasp the details of all variants.
• 5.4: This is again a very hand-wavy claim without actually being theoretical or empirical support.

Throughout the text, there are many inaccurate or somewhat sloppy statements and references that confuse a specialist in flow matching models (and would likely impede understanding by non-specialists as well). A list follows.

• Abstract: I would suggest to revise it to explain the problem setting and approach at a higher level.
  • First sentences: "Flow-based generative modeling is a powerful tool for solving inverse problems in physical sciences" -- this is a bold claim. The use of flow-based models for inverse problems is not yet well-established; in fact, this is what this paper aims to do.
  • The next few sentences do not set up the problem well (it is not even explained that we are talking about continuous normalising flows).
  • The results at the end do not make sense without context: what does "improves the accuracy by 53%" mean?
• Introduction:
  • L046: "normalising flows transform a noise distribution to the posterior distribution" is true, but:
    • It is a statement with low specificity (VAEs and GANs also transform noise to data).
    • These models are first introduced in the setting of training from samples, which is not what we usually have in Bayesian inference (no ground truth samples from the posterior).
    • Two of the three citations about diffusion models are actually about ODEs / flow matching. The two are of course connected, but I think it is unfair to call flow-based models an instance of "success of diffusion models [...] specifying a corruption process". For instance, flow ODEs can be learned that are not the probability flow ODEs of diffusion processes, including Liu et al.'s rectified flow (the first iteration is indeed a diffusion ODE, the later 'straightened' ones are not) and Tong et al.'s minibatch OT-based flow matching.
      • One solution could be to explicitly state the connection between FM and diffusion in the cases where it exists (e.g., for Ornstein-Uhlenbeck noise, optimal drifts of ODE and SDE are both expressed in terms of the score, so learning one is tantamount to learning the other).
  • In the paragraph starting L052, somehow we have jumped from a general distribution-matching setting (which is not how flow-based models were introduced -- they are trained from samples!) to a conditional posterior modelling setting. Please state the setting/assumptions (e.g., that we have conditional posterior samples from a simulator).
• Related work:
  • "Inverse problems with diffusion models": This seems to be better named "solving inverse problems under a diffusion model prior". Although the works that do this with Monte Carlo (e.g., Cardoso et al, Dou et al,) are mentioned (you could also consider Song et al.), there is also stochastic optimisation (e.g., Mardani et al., Graikos et al.), or amortisation by RL methods (e.g., Black et al., Fan et al., Venkatraman et al.).
  • "Flow matching": It is strange to see "optimal transport paths (Lipman et al.)" contrasted with "independent couplings or rectified flows (Liu et al., Tong et al.)". In fact, it is Rectified Flow and OT-CFM (Liu et al., Tong et al.) who consider non-independent couplings through rectification steps or OT couplings (thus actually approximating the dynamic OT), respectively, while Lipman et al.'s flow matching is equivalent to one using independent couplings and is OT only on the level of the conditional probability paths used for training.
• FM theory:
  • Equation (1): $\theta$ in subscript should be $\phi$.
  • L156: Because smoothness conditions are stated, they should be precise: Do you assume Bochner integrability? Continuous differentiability (how many times?) in both $\theta$ and $t$? It should also be stated that $p_t(x)=p(t,x)$ and $p$ is a function $[0,1]\times\mathbb{R}^d\to\mathbb{R}$.
  • L182 is hard to understand: what is meant by "$q(z)=p_1(\theta)$? It should be said that the conditioning variable $z$ is identified with the endpoint $\theta_1$, etc.
• Controls for improved accuracy:
  • LL201-203 do not make sense to me. The paragraph begins with conditioning of ODEs -- how is an old trick in diffusion "for example" w.r.t. such conditioning?
  • NB. Equation (7) will be the exact $t=1$ endpoint of integration ($\theta_1=\hat{\theta}_1$) if we have a perfectly fit OT or any model with straight integral curves (such as the converged ODE after many iterations of rectified flow)!
  • LL225-227 and later at LL352-355: Once again, I am surprised by the repeated discussion of Liu et al. and Tong et al. yet the omission of the actual algorithms they propose (rectification and MBOT coupling), which both produce straighter paths than a vanilla FM (and hence inference in fewer steps).

Experiment results are not convincing:

• Results are not consistently showing improvement (cf. Table 1) and error bars are not reported, making it impossible to assess significance.
• In Section 5, can you comment on computational efficiency in terms of wall time?
• Lensing:
  • The evaluations do not seem to guarantee coverage ($\chi^2$ is obviously not sufficient for this, and the SBC tests use projection only on a single parameter). Have you considered coverage tests such as those used in Legin et al., Section 3? Currently it is not demonstrated that the proposed method achieves more accurate posterior sampling.