Flow Matching for Posterior Inference with Simulator Feedback
generative modeling;simulation-based inference;astronomy
Yes
No ethics review needed.
Please see above.

Simpler than the encoder architecture, did you consider physics-inspired ways of providing the simulator information to the drift model? For example, simply expressing the drift as (learned vector field NN) + (learned scalar or diagonal NN) x (simulator gradient), as often done in work on diffusion models for sampling of Boltzmann distributions.
- The problem studied is highly relevant as foundation models (including diffusion and flow-based models) become available in various scientific domains. It is important to develop  inference methods for inverse problems that have low bias, good posterior coverage, and high efficiency of training and inferences -- this paper attempts to solve these problems.
- The proposed algorithm plausibly attacks these challenges (even if it is not very well demonstrated by the experiments, see below) and should give an asymptotically correct solution to the posterior sampling problem.
- Interesting analysis of algorithm variants in Section 5.2.
The paper studies the problem of modelling the posterior $p(\theta\mid x)$ in a generative model $\theta\to x$, where $p(x\mid\theta)$ is available as a simulator, but possibly without access to exact likelihoods. When the sampling of $\theta$ given $x$ is modelled as a conditional (on $x$) neural ODE and this ODE is trained by flow matching objectives, it is proposed to place an inductive bias on the drift model: the output of simulator or its gradient, evaluated at an intermediate time point or its extrapolation to the target space, is encoded and given as an input to the drift model. Such a form of the drift model is hypothesised to improve the approximation of the target distribution by effectively guiding the drift to the modes of the posterior. Experiments are done on several low-dimensional simulation-based inference tasks, including the lens and source parameter estimation problem in strong gravitational lensing.
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.](https://proceedings.mlr.press/v202/song23k.html)), there is also stochastic optimisation (e.g., [Mardani et al.](https://arxiv.org/abs/2305.04391), [Graikos et al.](https://arxiv.org/abs/2206.09012)), or amortisation by RL methods (e.g., [Black et al.](https://arxiv.org/abs/2305.13301), [Fan et al.](https://arxiv.org/abs/2305.16381), [Venkatraman et al.](https://arxiv.org/abs/2405.20971)).
  - "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.
Yes
No ethics review needed.
I would be happy if authors could provide arguments against my concerns.
- *Novelty*: To the best of my knowledge, the idea of including simulator feedback in flow matching for posterior estimation is novel
- *Soundness*: The idea of using simulator feedback to correct for the imprecision of score matching may indeed be helpful in certain applications. 
- *Presentation*: I found the tables and figure easy to read and informative, while being also visually appealing.
This paper introduces flow matching with simulator feedback for simulation-based inference which extends offline flow matching for neural posterior estimation with an online phase that uses online simulations to improve the accuracy of the estimated posterior distribution. Indeed, the authors observe that learning perfectly the score function corresponding to the true posterior distribution is challenging and propose to access the simulator online and correct these imperfections at evaluation time. The method is supposedly much more efficient than alternative online method such as MCMC while having the potential to provide as accurate results. The paper introduces two types of control signals, a gradient-based and a learning-based control signal, which are made for differentiable (and deterministic) and non-differentiable (and potentially stochastic) simulators respectively. The method is empirically tested on 4 common SBI benchmarking tasks and a "strong gravitational lensing" problem. Results highlight that simulator feedback help improving the accuracy of the posterior distributions.
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. 


[1]:http://proceedings.mlr.press/v130/lueckmann21a/lueckmann21a.pdf
Yes
No ethics review needed.
I would like to see especially the experiments and responses to the issues mentioned as weaknesses.
The paper has some strengths overall, which I will outline below. 


**Strengths:**
1. Introduction of the flow matching perspective for these set of problems making a faster inference procedure.
2. Providing the solutions on using either differentiable or non-differentiable simulators.
The paper considers the problem of solving inverse problems mostly in physical sciences. In particular, we are interested in posterior inference. This paper propose to use the flow matching perspective refined with additional control signals coming from a simulator. Moreover, the authors consider various scenarios depending on the differentiability of the simulator. Finally, they show the empirical results on a few simulator-based inference problems and the gravitational lensing inverse problem.
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.
Yes
No ethics review needed.
In addition to the weaknesses above: 
* What are the challenges to scaling the approach to high-dimensional spaces?
* How sensitive is the method to the quality of the simulator? What happens when the simulator contains significant approximations or errors?
* Have you explored using multiple different types of control signals simultaneously? Could this provide complementary benefits?
* Could the method be extended to handle multiple observations simultaneously in a more efficient way?
* The paper tackles an interesting problem of incorporating simulator feedback into generative models trained with flow matching. With the recent interest in the flow matching for learning generative models, the problem of incorporating downstream rewards, e.g. simulator feedback, is a critical one. This paper presents one of the first efforts in that direction (based on my knowledge). 
* The application to the lensing problem is quite interesting. Generative models are well suited to such scientific inverse problems, and this method could be a useful addition to the toolbox for such problems. 
* The method supports feedback from black-box simulators as well as differentiable simulators where gradient feedback is available. 
* The method also provides considerable speed-ups in the simulation over other approaches. 
* The paper is also quite clearly written and easy to follow.
This paper introduces a method to improve flow-based generative models for simulation-based inference by incorporating simulator feedback through control signals. The key idea is to refine a pretrained flow network with additional control signals based on simulator outputs, which can include gradients and problem-specific cost functions for differentiable simulators or learned signals from non-differentiable simulators. The authors demonstrate their method on several benchmark problems and show substantial improvements in accuracy (53%) while maintaining fast inference times (up to 67x faster than MCMC) on a challenging strong gravitational lensing application.
* 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.
Yes
No ethics review needed.
- Is there something particular about the flow matching setup (e.g. linear trajectories) that make simulator control applicable here, in contrast to a more traditional method like NPE? Does the method rely on assuming optimal transport coupling paths (eqs. 5-6), or is it more generally applicable?
- If I understand correctly, the goal of simulator control is to produce flow matching trajectories that better model the joint $(\theta, x_0)$ space by giving an additional signal that the interpolated parameter-generated samples should produce simulations consistent with the original $x_0$. This reduces the effect of simulator stochasticity. Could this partly be recreated by generating multiple samples from the same parameter point partway through training? Similarly, could producing a batch of samples for comparison to the $x_0$ further improve the flow matching control signal?
- In the strong lensing section, why was the comparison primarily made to HMC, rather than e.g. NPE?
- How critical was the $t > 0.8$ empirical threshold choice for controls, and do the authors expect this to be problem-dependent or fairly universal?
- Timely paper on an important problem in the field -- using information from the simulator to speed up simulation-based inference, in order to make it more sample-efficiency and well-calibrated.
- Discussion of several control mechanisms, which could be applicable depending on the specific simulator and domain.
- Good discussion of underlying theory, both the flow matching aspect as well as simulation-based inference. Sound comparison to surrounding context literature, e.g. relevance to classifier-free guidance.
- Application to an important and challenging problem in cosmology -- strong lens modeling, and comparison with a baseline HMC approach.
The paper presents a method for including simulator guidance in a flow matching-based simulation-based inference setup. The basic idea as far as I understand is to model the posterior using generative flow matching, and during training (possibly only at a fine-tuning stage), draw samples from the simulator by interpolating parameter samples from the velocity trajectory, and include an additional signal ensuring consistency between those samples and the training datapoint. The overall motivation is to increase simulation efficiency when a stochastic simulator is used for simulation inference (as is usually the case).
- 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.
We show that feedback from simulators gives a greatly improved posterior accuracy when solving inverse problems with flow matching
@inproceedings{
anonymous2024flow,
title={Flow Matching for Posterior Inference with Simulator Feedback},
author={Anonymous},
booktitle={Submitted to The Thirteenth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=DoDNJdDntB},
note={under review}
}
Flow-based generative modeling is a powerful tool for solving inverse problems in physical sciences that can be used for sampling and likelihood evaluation with much lower inference times than traditional methods. We propose to refine flows with additional control signals based on a simulator. Control signals can include gradients and a problem-specific cost function if the simulator is differentiable, or they can be fully learned from the simulator output. 
In our proposed method, we pretrain the flow network and include feedback from the simulator exclusively for finetuning, therefore requiring only a small amount of additional parameters and compute. We motivate our design choices on several benchmark problems for simulation-based inference and evaluate flow matching with simulator feedback against classical MCMC methods for modeling strong gravitational lens systems, a challenging inverse problem in astronomy. We demonstrate that including feedback from the simulator improves the accuracy by $53$%, making it competitive with traditional techniques while being up to 67x faster for inference. Upon acceptance, we will make our code publicly available.
I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
generative modeling simulation-based inference astronomy
I certify that there is no acknowledgement section in this submission for double blind review.
/pdf/e4189acdd8cbf40d196cb03999068c0c20bd07b9.pdf
applications to physical sciences (physics, chemistry, biology, etc.)
I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.
I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
Flow Matching for Posterior Inference with Simulator Feedback
ICLR 2025 Conference Submission
ICLR.cc/2025/Conference/Submission