Flow Matching for Scalable Simulation-Based Inference

Neural posterior estimation methods based on discrete normalizing flows have become established tools for simulation-based inference (SBI), but scaling them to high-dimensional problems can be challenging. Building on recent advances in generative modeling, we here present flow matching posterior estimation (FMPE), a technique for SBI using continuous normalizing flows. Like diffusion models, and in contrast to discrete flows, flow matching allows for unconstrained architectures, providing enhanced flexibility for complex data modalities. Flow matching, therefore, enables exact density evaluation, fast training, and seamless scalability to large architectures--making it ideal for SBI. We show that FMPE achieves competitive performance on an established SBI benchmark, and then demonstrate its improved scalability on a challenging scientific problem: for gravitational-wave inference, FMPE outperforms methods based on comparable discrete flows, reducing training time by 30% with substantially improved accuracy. Our work underscores the potential of FMPE to enhance performance in challenging inference scenarios, thereby paving the way for more advanced applications to scientific problems.

Paper

Similar papers

Peer review

Reviewer VmG37/10 · confidence 3/52023-07-03

Summary

The paper introduces a new solution for Simulation Based Inference (SBI) by proposing Flow Matching Posterior Estimation (FMPE) as an alternative to Neural Posterior Estimation (NPE). The authors leverage the power of Flow Matching (FM) generative models as amortizors to enhance the accuracy and efficiency of SBI. The authors emphasize that FM is a well-suited approach for SBI due to its advantageous characteristics, supported by empirical and theoretical evidence. These characteristics include easy sampling, tractable likelihood estimation, and mass coverage of the posterior distribution. The paper details a series of experiments conducted on domain-relevant tasks to demonstrate the superiority of the FMPE method. The results showcase the effectiveness of their proposed approach compared to existing methods.

Strengths

- Well written paper - Nice theoretical results which provide good intuitions despite constrained hypotheses

Weaknesses

- The likelihood of FM is very expensive (see Table 5 in App D) making it unsuited for SBI. This issue is not discussed in the paper. - The SBI benchmark suggests good sampling capabilities but Fig 4 in Sec 4 and in Fig 9 in App C suggest that the density estimation task is poorly performed (despite mode coverage being achieved). - The experimental results (see Fig 3 in Sec 4 for instance) don't convey a clear superiority of the proposed method over NPE.

Questions

- Can we have a reference for the computation of Eq (54) in App B ? (just like (48) was derived from (49)) - As mentioned in Sec 3.1, [1] provides a MSE bound of the KL in the SDE framework. Did you try such models ? - Can you further illustrate the superiority of your time prior ? - Can you comment on the poor histogram fit in Fig 4 in Sec 4 and in Fig 9 in App C ? Is this mismatch critical or can it be tolerated ? Can we see the histograms for NPE ? - How many integration steps were used in the experiments when sampling or computing the likelihood ? Was the dopri5 integrator always used ? - The computational cost of the FM's likelihood is very prohibitive (see Table 5 in App D) which makes the comparison in Table 1 a bit unfair. Do you have a solution to address this issue ? (Are you using the Hutchinson's trace estimator ?) - The training time difference between NPE and FMPE in GW (L294-L296) is very surprising. How do you explain it ? - What is the design of the Normalizing Flows for NPE experiments ? [1] Michael S Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.08797, 2023.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

Limitations

Not applicable.

Reviewer 7UFh6/10 · confidence 3/52023-07-06

Summary

This paper proposes solving Bayesian inference problems where we would like to estimate & sample from p(theta | x) when p(theta) and p(x | theta) are known. This is done by randomly sampling from the joint distribution and training conditional generative models as amortized approximate posteriors. Unlike ELBO-based training methods, the proposed approach does not ever sample from the model during training. Experiments on some benchmark problems seem to be done, though this reviewer has no experience in these problems.

Strengths

- Well-written, motivated, and easy to understand. - Method is straightforward but seems to work well. - Theorem 1 is great.

Weaknesses

- Benchmarks do not seem to be be challenging.

Questions

- How is JSD computed? This would typically require sampling from the true posterior. - Are the true posteriors known for these problems? - The inference problems may not be challenging enough. Typically, to fit a approximate posterior, we would use the KL(p_model || p_true_posterior) or equivalently, maximizing the ELBO for some given x. This is done because this allows us to sample from the approximate posterior, which has much better properties than sampling from the prior p(theta). I wonder if the reason this is not discussed is because the problems are just not very challenging. For instance, can the proposed approach be used to train a Bayesian neural network? - In regards to Theorem 1, have you tried estimating how large the value of C is for the experimental problems? This grows exponentially with respect to time right? - In regards to Theorem 1, in the author's opinions, what is the most problematic assumption that may not hold in practice?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

I do not see a discussion on limitations.

Authorsrebuttal2023-08-21

We thank the reviewer once more for their positive and helpful review. We'll be happy to further clarify if there are any questions before the discussion window closes at 1pm EDT today.

Reviewer SM255/10 · confidence 4/52023-07-07

Summary

The paper proposes a novel method for posterior estimation when samples of likelihood are given. Specifically, the authors propose the use of flow matching (FM) to train continuous normalizing flow-based (CNF-based) posterior distributions, a method they term "flow matching posterior estimation" (FMPE). The authors highlight that their proposed method is appealing for two primary reasons. Firstly, they emphasize that CNFs are more powerful than the conventional discrete normalizing flow-based (DNF-based) methods, as CNFs can utilize more flexible network architecture to represent complex densities. Secondly, FM eliminates the need to run integrators for training CNFs, a factor often dissuading researchers from using them. To train a CNF, one must define parametric vector fields (VFs) and train the model by maximizing the likelihoods, a process that includes running integrators. However, FM differs from the conventional maximum likelihood estimation (MLE) method in that it strives to match the model's VFs to predefined target VFs, from a source noise to a target distribution. The source noise, often referred to as the initial distribution, commonly uses the standard normal distribution. The target VFs are inferred by assuming that conditional distributions from a single sample of the target to the source distribution take Gaussian forms, in a manner similar to the forward process in diffusion-based generative models. Interestingly, one can sample the target VFs at any specified time between the source and the target distribution, circumventing the need for running integrators. In addition to the adoption of FM in posterior estimation, the paper demonstrates that under certain regularity assumptions, the flow matching loss provides an upper bound on the forward KL divergence between the target and the model posterior. The authors assert that this theoretical finding is crucial for posterior estimation as it suggests that the model posterior aims at covering the support of the target distribution rather than concentrating on higher density samples within the target's high probability region. The authors also introduce some techniques to enhance FMPE. For instance, they propose the use of pyramid-like architectures from data to latent, with gated linear units to incorporate (latent, time) dependence, in place of the commonly chosen U-Net-based network architectures for CNFs in generative models. Additionally, they suggest an alternative t-weighting in the loss, superseding the uniform weighting of the original FM. Lastly, the paper presents several experiments designed to demonstrate the effectiveness of the proposed method, including gravitational-wave inference tasks.

Strengths

Overall I find that the writing is clear, concise, and well-structured, making it easy for readers to follow the arguments and understand the key points. The authors have succeeded in providing a fresh perspective on the topic, shedding new light on the subject matter and offering meaningful contributions to the machine learning communities.

Weaknesses

From my perspective, the key innovation of this paper lies in its use of flow matching (FM) for posterior estimations. Essentially, the authors are training conditional models using FM, which marks a fresh approach in the field. While the aim of the study is clear, however, the technical challenges it faces seem minor. The authors could benefit from providing more details about these challenges to strengthen their argument and increase the depth of their research. The choice of experiments further aggravates these issues. For instance, the use of gravitational-wave inference tasks isn't compelling enough. To better showcase the technical contributions of their work, such as the new network architectures or the t-weighting in the loss, the authors should consider more demanding tasks. These could provide a more substantial platform for discussing their technical contributions.

Questions

N/A

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

Please refer to the comments provided in the weaknesses section.

Authorsrebuttal2023-08-21

We thank the reviewer once more for their positive and helpful review. We'll be happy to further clarify if there are any questions before the discussion window closes at 1pm EDT today.

Reviewer 1iUK7/10 · confidence 4/52023-07-07

Summary

This a **very well written** paper presenting an extension to the flow matching approach for generative modelling to the simulation-based inference setting. It goes in line with preceding works that adapted normalizing flows to SBI as well. The core idea of the paper is described in Equation (7): one can approximate (in average) the target posterior distribution by sampling $(\theta, x)$ from the joint pdf. This is rather unsurprising and it is how amortized neural posterior estimation (NPE) works. Nevertheless, this had not yet been done in the literature (at least to my knowledge) so it deserves to be shared to the community. Furthermore, the authors have the merit of including a section on the limitations of the posterior approximation (Section 3.1) which is not often done in the SBI literature. The experiments using the `sbi-benchmark` are a very nice addition as well. If I was going to mention one thing that bothered me about the paper was the fact of advertising their FMPE method as allowing for scalable SBI. In fact, what they present deals very well with large dimensions for the **observation space** (e.g. $x$) but the challenges related to high-dimensional **parameter space** (i.e. $\theta$) remains. In all, I was very pleased to read this paper and it pushed me to learn some new topics that I had only heard so far but had not yet had the time to dig further (flow matching, neural ODE, continuous normalizing flows). I'm voting for a 7 (accept) and am willing to increase my score if the authors answer my questions and consider my remarks/suggestions.

Strengths

- Very good writing and style of explanations. - Simple (in the sense of not being very far reached from previous works) yet relevant new method for simulation-based inference. - Validation experiments on the `sbi-benchmark` allowing for easy comparisons against previous works in the literature.

Weaknesses

- Results on real data seem rather preliminary and just for illustrative purposes. Previous work had already shown that GNPE was good for the gravitational waves setting because of its group invariances. Why one would want to use FMPE for it then? - Some mathematical notation are slightly inconsistent and hard to follow (I will be more precise in my questions) - Saying that the method is **scalable** without being very precise in which sense -- in observation space, not parameter space.

Questions

- Line 76: Is training time really one of the main goals when comparing methods for SBI? I would have expected more reflections around simulation budget (which, of course, is directly related to time). - I find it very strange to refer to $\theta$ as **samples** such as in Line 100. The authors also say "sample space" when often in the SBI literature we say rather "parameter space". - It was not super clear to me why you have to sample $\theta_1$ and what does this mean in the whole setup - Can you give more intuition in to the failure cases for the assumptions in Theorem 1? - By the end of your Section 3 there's no mention of how the choices of loss, architecture, and training procedure can help in improving the scalability of SBI. In fact, this feature is mainly mentioned in the title, abstract, and intro, but not that much commented in the rest of the paper. - Figure 3: You don't give any details about how the NPE was setup in the experiments. Can you give any intuition as to why things work so badly for the Lotka-Volterra example? - Line 296: Do you have any intuition as to why FPME trains so much faster as compared to the other methods considered in your experiments? - Line 313: What do you mean with the much simpler vector field of FMPE? Can you give more elements into this comment?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

The authors make no comment regarding the calibration and consistency of the posterior distribution obtained with their FMPE. Also, the results on the Lotka-Volterra in Figure 3 give the impression that the method can be improved for time series.

Area Chair oqmm2023-08-15

Author-reviewer discussion

Dear all, The author-reviewer discussion period has now started. It will continue for one more week, until August 21. @authors: Please respond to the comments or questions reviewers may further have. Remain short and to the point. @reviewers: Please read the author's responses and ask any further questions you may have. To facilitate the decision by the end of the process, please also acknowledge that you have read the responses and indicate whether you want to update your evaluation. - You can update your evaluation positively (if you are satisfied with the responses) or negatively (if you are not satisfied with the responses or share other reviewers' concerns). Please note that major changes are a reason for rejection. - You can also keep your evaluation unchanged. In this case, please indicate that you have read the responses and that you do not have any further comments. Best regards, The AC

Reviewer 1iUK2023-08-16

Thank you.

I thank the authors for their answers and clarifications. I did enjoy reading their paper and would be pleased to see them present it in December. However, I don't see specific reasons to increase my score even further, so I will just keep it as it is. Cheers, 1iUK

Reviewer VmG32023-08-18

First of all, I want to apologize for my very late reply to your rebuttal. Your answers to our reviews were really well made. Based on your answer to my review I fully understand that the actual bottleneck is the training process (despite the simple regressive objective) and likelihood computations in the inference phase are still negligible. I truly appreciate your new plot to investigate the added value of the power low as well as the experiment with the approximate divergence for GW inference. Like the other reviewers, I was also skeptical about the apparent simplicity of your experiments but I know fully understand that the toy experiments were meant to be sanity checks whereas the GW experiment is the real deal. Given this excellent rebuttal and the other reviews, I am ready to increase my score to a solid 7.

Program Chairsdecision2023-09-21

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC