Flow Matching for Scalable Simulation-Based Inference
Authors: Jonas Bernhard Wildberger
Maximilian Dax
Simon Buchholz
Stephen R Green
Jakob H. Macke
Bernhard Schölkopf
Venue: NeurIPS 2023 poster (2023)
Keywords: simulation-based inference
likelihood-free inference
machine learning for physical sciences

Abstract: 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.

Flow Matching for Scalable Simulation- Inference
Jonas Wildberger'
Maximilian Dax*
Max Planck Institute for Intelligent Systems Tübingen, Germany
Max Planck Institute for Intelligent Systems Tübingen, Germany
wildberger wwldberger.jonaseeeeee. jonssttuebnnennnnnn. mpg.en..pp...o de
maximilian www.:aandddtttueiinnnn.... mpg.. de
Simon Buchholz*
Stephen R. Green
Max Planck Institute for Intelligent Systems Tübingen, Germany
University of Nottingham Nottingham, United Kingdom
sbuchholz@tue. mpg. de
Jakob H. Macke
Bernhard Schölkopf Max Planck Institute for Intelligent Systems Tübingen, Germany
Max Planck Institute for Intelligent Systems &
Machine Learning in Science,
University of Tübingen Tübingen, Germany
Abstract
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 architec- tures, 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 com- petitive performance on an established SBI benchmark, and then demonstrate its improved scalability on challenging scientific problem: for gravitational-want in- ference, 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.
Introduction
The ability to readily represent Bayesian posteriors of arbitrary complexity using neural networks would herald revolution in scientific data analysis. Such networks could be trained using simulated data and used for amortized inference across Doservations--ringe myriad of scientific models. Thanks to innovative architectures such as normalizing flows [I] 2), approaches to neural simulation-barl based inference (SBI) 3 have seen remarkable progress in recent years. Here, we show that modern approaches to deep generative modeling (particularly flow matching) deliver substantial improvements in simplicity, flexibility and scaling when adapted to SBI.
tractable inference and speed to
EEqual contribution
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

g(
45th (0, x0)
(00m))
NPE
20 (0,0) x²)
vst (0,22)
0
1o.
18s, 1
900 x()
FMPE
Dt.z (0,)
Comparison of network architectures (left) and flow trajectories (right). Discrete flows (NPE, top) require specialized architecture for the density estimator. Continuous flows (FMPE, bottom) are based on a vector field parametrized with an unconstrained architecture. FMPE uses this additional flexibility to put an enhanced emphasis on the conditioning data which in the SBI context is typically high dimensional and in a complex domain. Further, the optimal transport path produces simple flow trajectories from the base distribution (inset) to the target.
Figure 1:
The Bayesian approach to data analysis is to compare observations to models the posterior distribution p(0)xr). (r)). This gives our degree of belief that model parameters 0 gave rise to observation x, and is proportional to the model likelihood p(z(0) times the prior p(0). One is typically interested in representing the posterior in terms of collection of samples. however obtaining these through standard likelihood-basee algorithms can be challenging for intractable or expensive likelihoods. In such cases, SBI offers alternative based instead data simulations x ~ p(=(0). Combined with deep generative modeling, SBI becomes a powerful paradigm for scientific inference [3. Neural posterior estimation (NPE) 4 61, for instance, trains conditional density estimator ((f[x: to approximate the posterior, allowing for rapid sampling and density estimation for any x consistent with the training distribution.
The NPE density estimator g(f|z) commonly taken to be a (discrete) normalizing flow [| 2], an approach that has brought state-of-the-ee performance in challenging problems such as gravitational- wave inference 77]. Naturally, performance hinges on the expressiveness of (0)zz). Normalizing flows transform noise samples through discrete sequence of basic transforms. These have been carefully engineered to be invertible with simple Jacobian determinant, enabling efficient maximum likelihood training, while at the same time producing expressive ((f|z)) ()). Although many such discrete flows are universal density approximators [2], in practice, they can be challenging to scale to very large networks, which needed for big-data experiments.
Recent studies g propose neural posterior score estimation (NPSE), rather different approach that models the posterior distribution with score-matching (or diffusion) networks. These techniques were originally developed for generative modeling [10) achieving state-of-the-art results many domains, including image generation [4]. Like discrete normalizing flows, difffusion models transform noise into samples, but with trajectories parametrized by a continuous "time" parameter t. The trajectories solve stochastic differential equation [5]] (SDE) defined in terms of vector field us which is trained to match the score of the intermediate distributions Pt. NPSE has several advantages compared to NPE, including the ability combine multiple observations at inference time [ and, importantly, the freedom to use unconstrained network architectures.
We here propose to use flow matching, another recent technique for generative modeling, for Bayesian inference, an approach we refer to as now-matching posterior estimation (FMPH). Flow matching also based on vector field i and thereby also admits flexible network architectures (Fig. I For flow

---

matching, however, directly defines the velocity field of sample trajectories, which solve ordinary differential equations (ODEs) and are deterministic. As consequence, flow matching allows for additional freedom in designing non-diffusion paths such as optimal transport, and provides direct
access to the density 16. These differences are summarized in Tab. [
NPE
NPSE
FMPE (Ours)
Tractable posterior density Unconstrained network architecture Network passes for sampling
Yes
No
Yes
No
Yes
Yes
Single
Many
Many
Table I: Comparison of posterior-estimation methods.
Our contributions are as follows:
We adapt flow-matching Bayesian inference, proposing FMPE. In general, the modeling requirements of SBI are different from generative modeling. For the latter, sample quality is critical, i.e., that samples lie in the support of complex distribution (e.g.. images). In contrast, for SBI, p(0 (z) is typically less complex for fixed æ, but x itself can be complex and digh-diiensional.. We therefore consider pyramid-Iis like architectures from z to Vt with gated linear units to incorporate (0, tt) dependence, rather than the typical U-Net used for images (Fig. T. We also propose an alternative t-weighting in the loss, which improves performance in many tasks.
Under certain regularity assumptions, we prove an upper bound on the KL divergence between the model and target posterior. This implies that estimated posteriors mass- covering, i.c., that their support includes all / consistent with observed æ, which is highly desirable for scientific applications [17].
We perform a number of experiments to investigate the performance of FMPE Our two- pronged approach, which involves a set of benchmark tests and real-world problem, is designed to probe complementary aspects of the method, covering breadth and depth of applications. First, on an established suite of SBI benchmarks, we show that FMPE performs comparably or better- than NPE across most tasks, and in particular exhibits mass- cass-covering posteriors all cases (Sec. 4. We then push the performance limits of FMPE on a challenging real-world problem by turning to gravitational- wave inference (Sec. [5) We show that FMPE substantially outperforms an NPE baseline terms of training time, posterior accuracy, and scaling to larger networks.
2 Preliminaries
Normalizing flows. A normalizing flow 2 defines probability distribution ((0lx)) over pa- rameters 0 E R² in terms of an invertible mapping the : I + IR" from simple base distribution (O).
((0[=) (v))).o(0) -o( (0)) det 00)
00
where (-). denotes the pushforward operator, and for generality we have conditioned on additional context r R.M. Unless otherwise specified, a normalizing flow refers discrete flow, where eiz is given by a composition of simpler mappings with triangular Jacobians, interspersed with shuffling of the 0. This construction results in expressive ((0l2r) and also efficient density evaluation 2
Continuous normalizing flows. A continuous flow [8] also maps from base to target distribution,
but is parametrized by a continuous "time" E [o, 1], where qo (0|z) (0) and qi (0)2)) (0)x)) For each f, the flow is defined by a vector field ur on the sample space velocity of the sample trajectories,
37 This corresponds to the
d dt (0) ve, Vr, (O)..
(9) 6.
(2)
Code availahle
In the SBI literature, this is also commonly referred to "parameter space".
3

---

We obtain the trajectories 8. = v0t,=(() by integrating this ODE. The final density is given by
q 2) = (V1..))-0(00) go exp (-fo- div ut (f))
(3)
which is obtained by solving the transport equation a.g. + div(gr - 0.
The advantage of the continuous flow is that (0) can be simply specified by a neural network taking R-+m+1 -> IR"", in which case (2) is referred to as a neural ODE [T8]. Since the density
is tractable via (3. it is in principle possible to train the flow by maximizing the (log-)likelihodd.
However, this often not feasible practice, since both sampling and density estimation require
many network passes to numerically solve the ODE (2).
Flow matching. An alternative training objective for continuous normalizing flows is provided by flow matching [16]. This directly regresses Ut,r on a vector field ue, that generates target probability path Pt.r- It has the advantage that training does not require integration of ODEs, however it is not immediately clear how to choose (ur, (t.,.) The key insight of [16] is that, if the path is chosen on sample-conditions basis I then the training objective becomes extremely simple. Indeed, given a sample-conditional probability path pr (010,) and a corresponding vector field uy (0(0))). we specify the sample-conditions flow matching loss as
LSCPM [0 a-p()) (L-p(0)a), 6 (..1)(9.101)) || (0,) (0.l000))
(4)
Remarkably, minimization of this loss is equivalent regressing (0) the marginal vector field ut, (0) that generates pa (e(0)z) [16]. Note that in this expression, the z-dependence of vt, (0) is
picked up via the expectation value, with the sample-conditional vector field independent of r.
There exists considerable freedom in choosing a sample-conditional path. Ref. [16] introduces the family of Gaussian paths
(5) where the time-dependent means (0) and standard deviations (81) can be freely specified (subject to boundary conditions b. For our experiments, we focus on the optimal transport paths defined by (0,) t0, and (6,)
(0(0)) = (V 14 (()
(1 mmin)t (also introduced in [16).. The sample- conditional vector field then has the simple form
ur (, = 01 (1 o min)
(6)
1
min))
Neural posterior estimation (NPE). NPE is an SBI method that directly fits density estimator ((0(z)) (usually& a normalizing flow) to the posterior p(0 x) 4 NPE trains with the maximum likelihood objective LNPE (p(0)p(219) log [f]x:)) using Bayes' theorem to simplify the expecta- tion value with (p(=)p(0)=)) + Ep(0hp(=10)- During training, LNPE is estimated based on an empirical distribution consisting of samples (f, 2) ~ p(f)ppzz()) Once trained, NPE can perform inference for every new observation using q(0)z)) thereby amortizing the computational cost of simulation and training across all observations. NPE further provides exact density evaluations of [f]zz) Both of these properties are crucial for the physics application in section 5 we aim to retain these properties with FMPE
Related work
Flow matching [16] has been developed as technique for generative modeling, and similar techniques are discussed in [19 21 and extended in [22 23). Flow matching encompasses the deterministic ODE special instance. Although our knowledge flow matching has not previously been applied to Bayesian inference, score- matching diffusion models have been proposed for SBI in 8 9 with impressive results. These studies, however, stochastic formulations via SDEs [15] or Langevin steps and are therefore not directly applicable when evaluations of the posterior density are desired (see Tab. T.. It should be noted that score modeling can also be used
version of diffusion models [10
*We refer to conditioning on as comple-conittoonng to distinguish from conditioning on r. 'The sample-come conditoona probability path should be chosen to be concentrated around 0; at t = (within a small region of size min) and be the base distribution at 0.

---

to parameterize continuous normalizing flows via an ODE. Extension of 8 9 to the deterministic formulation could thereby be seen as special case of flow matching, Many of our analyses and the practical guidance provided in Section therefore also apply to score matching.
We here focus on comparisons of FMPE against NPE 4 it best matches the requirements of the application in section [ Other SBI methods include approximate Bayesian computation [24-28), neural likelihood estimation [2 and neural ratio estimation [33/39] Many of these approaches have sequential versions, where the estimator networks are specifically tuned to a specific observation
Ta- FMPE has a tractable density, so it is straightforward to apply the sequential NPE 4 6 approaches
to FMPE. In this case, inference is no longer amortized, so we leave this extension to future work.
3 Flow matching posterior estimation
To apply flow matching SBI we use Bayes' theorem to make the usual replacement (p(()) a) + p(10) in the loss function 4 eliminating the intractable expectation values. This gives the FMPE loss
FEME Etsp()), 6 x1=p(()))=p(e(00)) (p-1)( (0.) (0(00(;))1.
(7)
which we minimize using empirical risk minimization over samples (0, p(0)p(z/00 In other words, training data is generated by sampling 9 from the prior, and then simulating data r correspond- ing to 0. This is in close analogy to NPF training, but replaces the log likelihood maximization with the sample-conditionll flow matching objective, Note that in this expression we also sample p(t), t E [0, 1] (see Sec. 3.3). which generalizes the uniform distribution in (4. This provides additional freedom to improve learning in our experiments.
3.1 Probability mass coverage
As we show in our examples, trained FMPE models q(0]zz) can achieve excellent results in approxi- mating the true posterior p(0]z). However, is not generally possible to achieve exact agreement due to limitations in training budget and network capacity. It is therefore important to understand how inaccuracies manifest. Whereas sample quality is the main criterion for generative modeling, for scientific applications one is often interested in the overall shape of the distribution. In particular, an important question is whether (((..)) )) is mass-covering, i.e., whether contains the full support of p(0|z)). This minimizes the risk to falsely rule out possible explanations of the data. It also allows us to use importance sampling the likelihood p(x[(0) of the forward model can be evaluated, which can be used for precise estimation of the posterior 40
Consider first the mass-covering property for NPE. NPE directly minimizes the forward KL divergence (p(0)p)))lll())))))
argmin, (gllp)
and thereby provides probability- mass covering results. Therefore, even if NPE is not accu- rately trained, the estimate ((((x)) should cover the entire support of the posterior p(f]zz) and the failure to do so can be observed in the validation loss. As an illustration in an unconditional setting, we observe that a unimodal Gaus- sian y fitted to bimodal target distribution p captures both modes when using the forward KL divergence DKL (pl Ig), but only a single mode when using the backwards direction (gl p)) (Fig,)
argmin, y DkL
argmin, LFM
Figure 2: Gaussian (blue) fitted to a bimodal distribution (gray) with various objectives.
For FMPE, we can fit a Gaussian flow-matching model (()) A(, g²) to the same bimodal target, in this case, parametrizing the vector field
(0) (af + (10)² 00 + fu az
(8)
L. (0 + (16)²)
we also obtain mass-covering distribution when fitting the learnable parameters (û, via ( This provides some indication that the flow matching objective induces mass-covering behavior, and leads us to investigate the more general question of whether the mean squared error
(see Appendix A
5

---

between vector fields th and by bounds the forward KL divergence. Indeed, the former agrees up to constant with the sample-conditional loss
(see Sec.
We denote the flows of u4, Ut, by &t, th, respectively, and we set 4: = (00))-00011 Pt = (0+) ) The precise question then is whether we bound DEL (Pa (1 by (S, w) 2 for some positive power a. It was already observed in [42] that this is not true in general, and we provide a simple example to that effect in Lemma] in Appendix B Indeed, it was found in [42 that to bound the forward KL divergence we also need to control the Fisher divergence, J [P (r(d)) (n): - (q:(0))²
Here show instead that bound can be obtained under sufficiently strong regularity assumptions on po. Ug. and Dt. The following statement is slightly informal, and we refer to the supplement for the complete version.
Theorem 1. Lei po go and assume and are two vector fields whose floseflows satisfy pl = (01)) ).po
and Qi = (11)-90- Assume that po is square integrable and satisfies AV po S c(1 + (0))
and un and v, have bounded second derivatives. Then there is a constant C > such that (for MSE,(u, v) < 1))
(PI(Pill I &CMSE, r)
The proof of this result can be found in appendix B
While the regularity assumptions are not guaranteed to hold in practice when v2 is parametrized by a neural net, the theorem nevertheless gives some indication that the flow-matching objective encourages mass coverage. In Section 4 5 this is complemented with extensive empirical evidence that flow matching indeed provides mass-covering estimates.
We remark that was shown in 43 that the KL divergence of SDE solutions can be bounded by the MSE of the estimated score function. Thus, the smoothing effect of the noise ensures mass coverage, an aspect that was further studied using the Bokker-Planck equation in For flow matching, imposing the regularity assumption plays a similar role.
3.2 Network architecture
Generative diffusion or flow matching models typically operate on complicated and high dimensional data in the 0 space (e.g., images with millions of pixels). One typically uses U-Net 44 like architectures, as they provide natural mapping from to a vector field (0) of the same dimension. The dependence on and an (optional) conditioning vector is then added on top of this architecture.
For SBI, the data x is often associated with a complicated domain, such as image or time series data, whereas parameters 0 are typically low dimensional. In this context, it is therefore useful to build the architecture starting
mapping from x to u(x:) and then add conditioning on 0 and l. In practice, one can therefore use any established feature extraction architecture for data in the domain of x and adjust the dimension of the feature vector to n dim(@). In our experiments, we found that the (t, 0) condiiiiin is best achieved using gated linear units [45) to the hidden layers of the network (see also Fig. DD: these are also commonly used for conditioning discrete flows on æ.
3.3 Re-scaling the time prior
The time prior (0, 1] in ④ distributes the training capacity uniformly across t. We observed that this not always optimal in practice, as the complexity of the vector field may depend on For FMPE we therefore sample t in 7 from power-law law distribution Pa (t) +!/(1+a), t e j0, 1]. introducing an additional hyperparameter a. This includes the uniform distribution for a = 0, but for a > 0. assigns greater importance to the vector field for larger values of (. We empirically found this to improve learning for distributions with sharp bounds (e.g., Two Moons in Section 4
4 SBI benchmark
We now evaluate FMPE on ten tasks included in the benchmark presented 46], ranging from
simple Gaussian toy models to more challenging SBI problems from epidemiology and ecology, with
varying dimensions for parameters (dim(0) E [2, 10) and observations (dim(æ)) E [2. 100). For
each task, we train three separate FMPE models with simulation budgets N {103, 10a 102). We

---

use simple network architecture consisting of fully connected residual blocks 47 to parameterize the conditional vector field. For the two tasks with dim(x) 100 (B-GLM-Rave SLCP-D), we
condition on (t, 0) via gated linear units as described in Section 3.2 the corresponding performance gain). For the remaining tasks with dim(
(Fig.) g in Appendix C shows
10 we concatenate
(t, 0, instead. We reserve 5% of the simulations for validation. See Appendix C for details.
For each task and simulation budget, we evaluate the model with the lowest validation loss by comparing q(0)az (r) to the reference posteriors p(0 xr) provided in 46] for ten different observations x in terms of the C2ST score [48 49 This performance metric is computed by training a classifier to discriminate inferred samples @ ~ g(0]2z) from reference samples 0 ~ p(f[zz). The C2ST score is then the test accuracy of this classifier, ranging from 0.5 (best) to 1.0. We observe that FMPE exhibits comparable performance to an NPE baseline model for most tasks and outperforms on several (Fig. 4. In terms of the MMD metric (Fig. 6 in the Appendix), FMPE clearly outperforms NPE (but MMD) can be sensitive to its hyperparameters [46). As NPE is one of the highest ranking methods for many tasks in the benchmark, these results show that FMPE indeed performs competitively with other existing SBI methods. We report additional baseline for score matching in Fig. 7 in the Appendix.
As NPE and FMPE both directly target the posterior with a density estimator (in contrast to most other SBI methods), observed differences can be primarily attributed to their different approaches for density estimation. Interestingly, a great performance improvement of FMPE over NPE is observed for SLCP with large simulation budget (N = 10"). The SLCP task is specifically designed to have a simple likelihood but a complex posterior, and the FMPE performance underscores the enhanced flexibility of the FMPE density estimator.
Finally, we empirically investigate the mass coverage suggested by our theoretical analysis in Section [.1 We display the density log (0(xr) of the reference samples 0 ~ p(0) x) under our FMPE model as a histogram (Fig. 3) All samples 0 p(0[zr) fall into the support from (()(z)). This becomes appar- ent when comparing to the density log q(0)z² x) for samples 0 ~ g(0lzz) x)) from itself. This FMPE re- sult is therefore mass covering. Note that this does not necessarily imply conservative posteriors (which also not generally true for the forward KL diver- gence [I7 50 (I and some parts of p(0[z:) may still be undersampled. Probability mass coverage, however, implies that part is entirely missed (com- pare Fig. 2. even for multimodal distributions such as Two Moons. Fig. 9 in the Appendix confirms the mass coverage for the other benchmark tasks.
0 p~0()) (~(()))))
2
2
log g00lz)
Histogram of FMPE densities log (())z) for reference samples 0 ~ p(f]x) (Two Moons task, N = 10a). The estimate q(0 (:) clearly covers p(0) (r) entirely.
Figure
5
Gravitational-wawy inference
5.1 Background
Gravitational waves (GWs) are ripples of spacetime predicted by Einstein and produced by cata- clysmic cosmic events such as the mergers of binary black holes (BBHs). GWs propagate across the universe to Earth, where the LIGO- LIGO-Virgo- Hirgo-KAGRA observatories measure faint time-series signals embedded in noise. To-date, roughly 90 detections of merging black holes and neutron stars have been made [52]. all of which have been characterized using Bayesian inference to compare against theoretical models [1 These have yielded insights into the origin and evolution of black holes (53) fundamental properties of matter and gravity [54] 55), and even the expansion rate of the universe [56].
Under reasonable assumptions on detector noise, the GW likelihood is tractable and inference is
'BBH parameters 0 E R 15 include black-hole masses, spins, and the spacetime location and orientation of the system (see Tab. 4 in the Appendix). We represent x in frequency domain; for two LIGO detectors and complex f [20, 512] Hz, Af = 0.125 Hz, we have E 16644
"Noise is assumed to be stationary and Gaussian, so for Diquency-domaile data, the GW likelihood p(x[0) = ((((())))))()) Here ls()) is theoretical signal model based on Einstein's 's theory of general relativity, and S., is the power spectral density of the detector noise.
7

---

1.0
GL
GL.-U
GM
Two Moons
0.8
NPE
SLCP
FMPF
0.6
C2ST
1.0
B-GLM
B-GLM-Raw
SIR
0.8
SLCP-D
L.
0.6
10a
10
10
10a
10
10
10a
10
10
10a
10
109
10a
10a
10a
Number of Simulations
Figure 4: Comparison of FMPE with NPE, standard SBI method, across 10 benchmark tasks 46
typically performed using tools (57 60]
based on Markov chain Monte Carlo 66 62] nested
sampling [63] algorithms. This can take from hours to months, depending on the nature of the event and the complexity of the signal model, with a typical analysis requiring up evaluations. The ever-increasing rate of detections means that these analysis times risk becoming a bottleneck. SBI offers promising solution for this challenge that has thus been actively studied in the literature [64H68, (Z] 69, 770 41]. A fully amortized NPE-based method called DINGO recently achieved accuracies comparable to stochastic samplers with inference times of less than minute per event [7]. To achieve accurate results, however, DINGO uses group-equivariant NPE [7] 69] (GNPE), an NPE extension that integrates known conditional symmetries. GNPE, therefore. does not provide tractable density, which is problematic when verifying and correcting inference results using importance sampling [4T)
10 likelihood
5.2 Experiments
We here apply FMPE to GW inference. As a baseline, we train an NPE network with the settings described in 7 with few minor changes (see Appendix D B This uses embedding network [71] to compress x to a 128-dimensional feature vector, which is then used to condition a neural spline flow 72 The embedding network consists of learnable linear layer initialized with principal components of GW simulations followed by series of dense residual blocks [47) This architecture powerful feature extractor for GW measurements (7]. As pointed out in Section 3.2] it is straightforward to reuse such architectures for FMPE, with the following three modifications: (1) we provide the conditioning on (t, () to the network via gated linear units in each hidden layer; (2) we change the dimension of the final feature vector to the dimension of 0 so that the network parameterizes the conditional vector field (t, x, 8) - (t)= (3) we increase the number and width of the hidden layers to use the capacity freed up by removing the discrete normalizing flow.
We train the NPE and FMPE networks with 5 10 simulations for 400 epochs using batch size of 4096 on an A100 GPU. The FMPE network (1.9 10 learnable parameters, training takes R 2 days) is larger than the NPE network (1.3 10 learnable parameters, training takes z: 3 days), but trains substantially faster. We evaluate both networks on GW150914 (73), the first detected GW. We generate reference posterior using the method described in 41]. Fig. compares the inferred posterior distributions qualitatively and quantitatively in terms of the Jensen- Sensen-Shannon divergence (ISD) to the reference
FMPE substantially outperforms NPE in terms of accuracy. with a mean JSD of 0.5 mnat (NPE: 3.6 mnat), and max JSD 2.0 mnat, an indistinguishabilityy criterion for GW posteriors [59]. Remarkably, FMPE accuracy is even comparable to GNPE, which leverages physical symmetries
*Our implementation builds on the public DINGO code from https: / /ttpttttttttut/////// com/ dingddgngngg/ooo/ gu/dingo "We omit the three parameters de, I, 0,n in the evaluation we phase marginalization in importance sampling and the reference therefore uses different basis for these parameters 41.. For GNPE we report the results from 77. which are generated with slightly different data conditioning. Therefore, we do not display the GNPE results in the corner plot, and the JSDs serve only as rough comparison. The JSD for the ta parameter is not reported in 77] due & marginalized reference.

---

NPE
FMPE
GNPE
Reference
NPE
FMPE
m1
1.2
1.3
0.8
2.5
1.1
JSD
(1
3.2
0.8
0.2
[mnat]
te [ms]
1.6
1.0
0.3
20
0.8
0.1
0.5
15
25
f2
0.4
0.3
0.5
11
0.3
0.2
0.1
10
dr
4.4
0.1
0.8
9.1
0.6
9.0
10.1
0.6
0.7
6
8.6
0.5
1.4
-65
00
1 0.6
0.1
0.2
d1. [Mpc]
4 ims]
Figure 5: Results for GW150914 (73]. Left: Corner plot showing marginals on the diagonal and 2D 50% credible regions. We display four GW parameters (distance dL, time of arrival te and sky coordinates a, d); these represent the least accurate NPE parameters. Right: Deviation between inferred posteriors and the reference, quantified by the Jensen-Shannon divergence (ISD). The FMPE posterior matches the reference more accurately than NPE, and performs similarly to symmetry- enhanced GNPE. (We do not display GNPE results on the left due to different data conditioning settings in available networks.)
to simplify data. Finally, find that the Bayesian evidences inferred with NPE (log p(z) -
and FMPE (log p(x) (1067.969++0.000 are consistent within their statistical uncertainties. correct evidence is only obtained in importance sampling when the inferred posterior ((0]z)) covers the entire posterior p(0 41. so this another indication that FMPE indeed induces mass-covering posteriors.
5.3
Discussion
Our results for GW150914 show that FMPE substantially outperforms NPE on this challenging problem. We believe that this is related the network structure follows. The NPE network allocates roughly two thirds of its parameters to the discrete normalizing flow and only one third to the embedding network (i.e., the feature extractor for x)) Since FMPE parameterizes just a vector field (rather than a collection of splines in the normalizing flow) it can devote its network capacity to the interpretation the nigh-dimensional achieves higher accuracy. Remarkably, the performance iscomparable to GNPE, which involves much simpler learning task with likelihood symmetries integrated by construction. This enhanced performance, comes in part at the cost of increased inference times, typically requiring hundreds of network forward passes. See Appendix D for further details.
IR 15744 Hence, scales better to larger networks and
In future work we plan to carry out a more complete analysis of GW inference using FMPE, Indeed, GW150914 14 is loud event with good data quality, where NPE already performs quite well. DINGO with GNPE has been validated in variety of settings 77 69 41 774 including events with larger performance gap between NPE and GNPE (69]. Since FMPE (like NPE) does not integrate physical symmetries, it would likely need further enhancements to fully compete with GNPE. This may require a symmetry-aware architecture 75), or simply further scaling to larger networks. straightforward application of the GNPE mechanism to -MPE-GFMPE-- also possible, but less practical due to the higher inference costs of FMPE. Nevertheless, our results demonstrate that FMPE a promising direction for future research in this field.

---

6 Conclusions
We introduced flow matching posterior estimation, a new simulation- inference technique based on continuous normalizing flows. In contrast to existing neural posterior estimation methods, it does not rely on restricted density estimation architectures such as discrete normalizing flows, and instead parametrizes distribution in terms of a conditional vector field. Besides enabling flexible path specifications, while maintaining direct access to the posterior density, we empirically found that regressing on a vector field rather than entire distribution improves the scalability of FMPE compared to existing approaches. Indeed, fewer parameters are needed to learn this vector field, allowing for larger networks, ultimately enabling to solve more complex problems. Furthermore, our architecture for FMPE (a straightforward ResNet with GLU conditioning) facilitates parallelization and allows for cheap forward/hackward passes.
We evaluated FMPE on a set 10 benchmark tasks and found competitive or better performance compared to other simulation-based inference methods. On the challenging task of gravitational wave inference, FMPE substantially outperformed comparable discrete flows, producing samples on par with method that explicitly leverages symmetries simplify training. Additionally, flow matching latent spaces are more naturally structured than those of discrete flows, particularly when using paths such optimal transport. Looking forward, it would be interesting to exploit such structure in designing learning algorithms. This performance and flexibilty underscores the capability of continuous normalizing flows to efficiently solve inverse problems.
Acknowledgements
We thank the DINGO team for helpful discussions and comments. We would like to particularly acknowledge the contributions of Alessandra Buonanno, Jonathan Gair, Nihar Gupte and Michael Pürrer. This material is based upon work supported by NSF's LIGO Laboratory which a major facility fully funded by the National Science Foundation. This research has made use of data or software obtained from the Gravitational Wave Open Science Center (gw-openscience.org)) a service of LIGO Laboratory, the LIGO Scientific Collaboration, the Virgo Collaboration, and KAGRA. LIGO Laboratory and Advanced LIGO are funded by the United States National Science Foundation (NSF) as well as the Science and Technology Facilities Council (STFC) of the United Kingdom, the Max-Planck-Sociey (MPS), and the State of Nicdersschheen/Germang for support of the construction of Advanced LIGO and construction and operation of the GEO600 detector. Additional support for Advanced LIGO was provided by the Australian Research Council. Virgo funded, through the European Gravitational Observatory (EGO), by the French Centre National de Recherche Scientifique (CNRS), the Italian Istituto Nazionale di Fisica Nucleare (INFN) and the Dutch Nikhef, with contribations by institutions from Belgium, Germany, Greece, Hungary, Ireland, Japan, Monaco, Poland, Portugal, Spain. The construction and operation of KAGRA are funded by Ministry of Education, Culture. Sports, Science and Technology (MEXT), and Japan Society for the Promotion of Science (JSPS), National Research Foundation (NRF) and Ministry Science and ICT (MSIT) in Korea, Academia Sinica (AS) and the Ministry of Science and Technology (MoST) in Taiwan. M.D. thanks the Hector Fellow Academy for support. J.H.M. and B.S. are members of the MLCOE, EXC number 2064/1- Project number 390727645 and the Tübingen AI Center funded by the German Ministry for Science and Education (FKZ 011S11033 180399))A).
References
[1] Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International Conference on Machine Learning. pages 1530-1538, 2015. arxiv: 15... 15770
[2] George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. Journal of Machine Learning Research, 22(57): 2021. URL: http:// wttpp/rrrrrg/pppp////222
19-1028.htm
[3] Kyle Cranmer, Johann Brehmer, and Gilles Louppe. The frontier of simulation-basel inference. Proc. Nat. Acad. Sci., 117(49):30055-30062, 2020. arxiv: .01429 doi doi:10.a 1073/
pnas 1912789117
10

---

[4] George Papamakarios and lain Murray. Fast E-free inference of simulation models with Bayesian conditional density estimation. In Advances in neural information processing systems, 2016.
arriiv: 1eXIvv1605..o 06376
[5] Jan-Matthis Lueckmann, Pedro J Gonçalves, Giacomo Bassetto, Kaan Öcal, Marcel Non- nenmacher, and Jakob H Macke. Flexible statistical inference for mechanistic models of neural dynamics. In Proceedings the 31st International Conference Neural Information Processing Systems, pages 1289-1299, 2017.
[6] David Greenberg, Marcel Nonnenmacher, and Jakob Macke. Automatic posterior transformation for likelihood-free inference. In International Conference on Machine Learning, pages 2404- 2414. PMLR, 2019.
[7] Maximilian Dax, Stephen R. Green, Jonathan Gair, Jakob H. Macke, Alessandra Buonanno, and Bernhard Schölkopf. Real-Time Gravitational Wave Science with Neural Posterior Estimation. Phys. Rev. Lett., 127(24):241003,, 2021. arXiv: 2106. 12594
doi: :oi:10... 1100 1103/PyysseLeett
241103
[8] Louis Sharrock, Jack Simons, Song Liu, and Mark Beaumont. Sequential neural score estima- tion: Likelihood free inference with conditional score based diffusion models. arXiv preprint arXiv:2210.04872, 2022.
[9] Tomas Geffner, George Papamakarios, and Andriy Mnih. Score modeling for simulation-based inference. arXiv preprini arXiv:2209. /4249, 2022.
[10] 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-2265. PMLR, 2015.
[11] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances neural information processing systems, 32, 2019.
[12] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-6855 6851, 2020.
[13] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 8780-8794. Curran Associates, Inc., 2021. URL: https/ttps//pprrcccciingg meu.iiinisi..nuuipppe nwwrips. 0011://h11/000044444444400000060
[14] Jonathan Ho, Chitwan Saharia, William Chan, David Fleet, Mohammad Norouzi, and Tim Salimans. Cascaded diffusion models for high fidelity image generation. J. Mach. Learn. Res., 23:47:1-47:33, 2022. URL: //jt:://mmr. http://jiwrroorrrpprereeeeee wrtgppapers/v2//21-00000 0635. html
[15] 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.1344 2020.
[16) Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. CoRR, abs/2210.02747, 2022. 10. 48550/arXiv.... 2210.02747
doi
[17] Joeri Hermans, Amaud Delaunoy, François Rozet, Antoine Wehenkel, and Gilles Louppe. Averting a crisis in simulation- inference. arXiv preprini 4rXiv:2110.06581, 2021.
[18] Tian Qi Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural or- dinary differential equations. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa-Bianchi, and Roman Garnett, editors, Advances in Neu- Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurlPS 2018, December 3-8, 2018, Montréal, Canada, pages 2572-6583, 2018. URL: http:::///::://poedddmma www.oeeeiinn.neurriirroooreoeeee neurips ec/ wwt..cccpppppeeee. (paper/208//aaeeeee 8/hash/ 69386f6bb htw.abbbe..ffaeeee22222222633333 Abstract.html

---

[19] Michael S Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants. arXiv preprint arXiv:2209.15571, 2022.
[20] Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint (7Xiv:2209.03000 2022.
[21] Kirill Neklyudov, Daniel Severo, and Alireza Makhzani. Action matching: A variational method for learning stochastic dynamics from samples. arXiv preprint arXiv:2210.06662, 2022.
[22] Aram Davtyan, Sepehr Sameni, and Paolo Favaro. Randomized conditional flow matching for video prediction. arXiv preprint vrXiv:221..4557 2022.
[23] Alexander Tong, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Brotor-Brooks, Kilian Fatras, Guy Wolf, and Yoshua Bengio. Conditional flow matching: Simulation free dynamic optimal transport. arXiv preprint arXiv:2302.00482, 2023.
[24] Scott A Sisson, Yanan Fan, and Mark A Beaumont. Overview of abc. In Handbook of approximate Bayesian computation, pages 3-54. Chapman and Hall/CRC, 2018.
[25] Mark A Beaumont, Wenyang Zhang, and David Balding. Approximate bayesian computation in population genetics. Genetics, 162(4))2025-2035, 2002.
[26] Mark A Beaumont, Jean-Marie Cornuet, Jean-Michel Marin, and Christian P Robert. Adaptive approximate bayesian computation. Biometrika, 96(4):983-990 2009.
[27] Michael G. B. Blum and Olivier François. Non-linear regression models for ap- proximate bayesian computation. Stat. Comput., 20(1):63-73, 2010. :01:10.10000 611222-009-9116-
[28] Dennis Prangle, Paul Fearnhead, Murray P. Cox, Patrick J. Biggs, and Nigel P. French. Semi- automatic selection of summary statistics for abe model choice, 2013. arxivv: 1302. 5624
[29] Simon Wood, Statistical inference for noisy nonlinear ecological dynamic systems. Nature, 466:1 102-4, 08 2010. doi : 10.11...18.00 10t.///
[30] Christopher C Drovandi, Clara Grazian, Kerrie Mengersen, and Christian Robert. Approximat- ing the likelihood in approximate bayesian computation, 2018 arxiiv 06645
[31] George Papamakarios, David Sterratt, and lain Murray. Sequential neural likelihood: Fast likelihood-free inference with autoregressive flows In The 22nd International Conference Artificial Intelligence and Statistics, pages 837-843 848. PMLR, 2019.
[32] Jan-Matthss Lueckmann, Giacomo Bassetto, Theofanis Karaletsos, and Jakob H Macke. Likelihood inference with emulator networks. In Symposium Advances in Approx- imate Bayesian Inference, pages 32-53. PMLR, 2019.
[33] Rafael Izbicki, Ann Lee, and Chad Schafer. High-dimensional density ratio estimation with extensions approximate likelihood computation. In Artificial intelligence and statistics, pages 420-429. PMLR, 2014.
[34] Kim Pham, David Nott, and Sanjay Chaudhuri. note on approximating abc-meme using flexible classifiers. Stat, 3, 03 2014. doi :00 1002/ stssstaa... 556
[35] Kyle Cranmer, Juan Pavez, and Gilles Louppe. Approximating likelihood ratios with calibrated discriminative classifiers. arXiv preprint arXiv: 02169, 2015.
[36] Joeri Hermans, Volodimir Begy, and Gilles Louppe. Likelihood-free meme with approximate likelihood ratios. arXiv preprint arXiv:/903.. 1et..999300577 10, 2019.
[37] Conor Durkan, Iain Murray, and George Papamakarios. On contrastive learning for likelihood- free inference. In International conference on machine learning, pages 2771-2781. PMLR, 2020.
[38] Owen Thomas, Ritabrata Dutta, Jukka Corander, Samuel Kaski, and Michael U. Gutmann. Likeiihood-fres inference by ratio estimation, 2020. arXiv:1611. 10242
12

---

[39] Benjamin K Miller, Christoph Weniger, and Patrick Forré. Contrastive neural ratio estimation. Advances Neural Information Processing Systems, 35:3262-3278 2022.
[40] Thomas Müller, Brian McWilliams, Fabrice Rousselle, Markus Gross, and Jan Novák. Neural importance sampling. ACM Transactions on Graphics (TOG), 38(5):1-19, 2019.
[41] Maximilian Dax, Stephen R. Green, Jonathan Gair, Michael Pürrer, Jonas Wildberger, Jakob H. Macke, Alessandra Buonanno, and Bernhard Schölkopf. Neural Importance Sampling for Rapid and Reliable Gravitational-Wav Inference. Phys. Rev. Lett. 130(17):171403, 2023. MrXiv:2214 05686
doi: 10. 1103/ 130.171403
[42] Michael S Albergo. Nicholas M Boffi, and Eric Gunden-Eijnded Stochastic interpolants: unifying framework for flows and diffusions. arXiv preprint arXiv:2303.08797, 2023.
[43] 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.
[44] Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical Image Computing and Computer-Assister Intervention-MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18, pages 234-241. Springer, 2015.
[45] Yann N Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated convolutional networks. In International conference on machine learning, pages 933-941, PMLR, 2017.
[46] Jan-Matthis Lueckmann, Jan Boelts, David Greenberg, Pedro Goncalves, and Jakob Macke. Benchmarking simulation- based inference. In International Conference Artificial Intelligence and Statistics, pages 343-351. PMLR, 2021.
[47] Kaiming He, Xiangyu Zhang. Shaoging Ren, and Jian Sun. Deep residual learning for image recognition, 2015. arXiv: 03385
[48] Jerome H Friedman. On multivariate goodness-of-fit and two-sample testing. Statistical Problems Particle Physics, Astrophysics, and Cosmology, 1:311, 2003.
[49] David Lopez- Paz and Maxime Oquab. Revisiting classifier two-sample tests. arXiv preprint arXiv:1610.06545, 2016.
[50] Arnaud Delaunoy, Joeri Hermans, François Rozet, Antoine Wehenkel, and Gilles Louppe. Towards reliable simulation- based inference with balanced neural ratio estimation, 2022, arXiv: 2208. 13624
[51] Arnaud Delaunoy, Benjamin Kurt Miller, Patrick Forré, Christoph Weniger, and Gilles Louppe. Balancing simulation-based inference for conservative posteriors. arXiv preprint arXiv:2304. 10978, 2023.
[52] R. Abbott et al. GWTC-3: Compact Binary Coalescences Observed by LIGO and Virgo During the Second Part of the Third Observing Run. arXiv preprint arXiv:2111.03606, 2021.
arXiv:2111 03606
[53] R. Abbott et al. Population Properties of Compact Objects from the Second LIGO- Gravitational- Wave Transient Catalog. Astrophys. J. Lett., 913(1):L7, 2021. arXiv 2010
14533 doi: 10. /w0.338811-8113//28822
[54] B. P. Abbott et al. GW1708111 Measurements of neutron star radii and equation of state. Phys. Rew. Lett., 121(16):11101, 2018. arXiv: 11581 doi: 10: 10::10.1103/PPpp.... 10t.101/Pyyssee..ttt..... 121 161101
[55] R. Abbott et al. Tests of general relativity with binary black holes from the second LIGO- gravitational-wave transient catalog. Phys. Rev. D, 103(12):122002, 2021. arXiv 2010 14529
10. 1103/Phssevv) 103. 122002
13

---

[56] B. P. Abbott et al. A gravitational standard siren measurement of the Hubble constant. Nature, 551(7678):85-88, 2017. arXiv: 1710. 05835 doi 1o1:10.1038/nattii www.//0ntntte24111
[57] Veitch, V. Raymond, B. Farr, W Farr, P. Graff, S. Vitale, et al. Parameter estimation for compact binaries with ground-based gravitational wave observations using the LALInference software library. Phys. Rev., D91(4))042003, 2015.
arXiv : 1409. 7215
doi: 10. 1103/
PhysNevD. 91. 042000 042003
[58] Gregory Ashton et al. BILBY: A user-friendly Bayesian inference library for gravitational- wave astronomy, Astrophys. J. Suppl., 241(2):27, 2019. 1811.00411..204 02042 doi : 3847/
(538-4366/ab0ffc
[59] I. M. Romero-Shaw et al. Bayesian inference for compact binary coalescences with bilby: validation and application to the first LIGO-Virgo gravitational-wave transient catalogue. Mon. Not. Roy. Astron. Soc., (99(3):3295-3311) 2020. 20xivv:00660077 1o0::/0.093//m staa2850
[60] Joshua S Speagle. dynesty: a dynamic nested sampling package for estimating Bayesian poste- riors and evidences, Monthly Notices of the Royal Astronomical Society, 493(3):3132-31558 Feb 2020. URL: http: //dx /ttpdxx.ioooooo/oooooooooo.doedt www.doi.oo..... org/ 10. 1093/mmras/staa277 arXiv: 1904. 02180
10. 1093/ htt3//mmass/sssee staa278
[61] Nicholas Metropolis, Arianna W Rosenbluth, Marshall N Rosenbluth, Augusta H Teller, and Edward Teller. Equation of state calculations by fast computing machines. The journal of chemical physics, 21(6): 1087-1022,,, 1953.
[62] W. K. Hastings. Monte Carlo sampling methods using Markov chains and their applica- tions. Biometrika, 57(1):97-109, 04 1970. httiv: http http::://aaaamcmccccco oup. com bin omet/
# pdf :o 1093/ biw.et/77. 1 1.97
[63] John Skilling. Nested sampling for general Bayesian computation. Bayesian Analysis, 1(4):833 859, 2006. doi: : 10. 0014/
[64] Elena Cuoco, Jade Powell, Marco Cavaglià, Kendall Ackley, Michal Bejger, Chayan Chatterjee, Michael Coughlin, Scott Coughlin, Paul Easter, Reed Essick, et al. Enhancing gravitational- wave science with machine learning. Machine Learning: Science and Technology, 2(1):01 1002,
5 2020.
arXiv: 03745
doi: 10. 2652222222223332332
[65] Hunter Gabbard, Chris Messenger, Ik Siong Heng, Francesco Tonolini, and Roderick Murray-Smithh Bayesian parameter estimation using conditional variational autoencoders for gravitational wave astronomy. Nature Phys., 18(1):112-117, 2022. arxiv: 06296
doi: 10. 1038/ 9415770011--11425
[66] Stephen R. Green, Christine Simpson, and Jonathan Gair. Gravitational-wave parameter estimation with autoregressive neural network flows. Phys. Rev. D, 102(10):104057, 2020.
arXiv: 07656 doi : 10. 100.1100//hhysee/ 102. 104057
[67] Arnaud Delaunoy, Antoine Wehenkel, Tanja Hinderer. Samaya Nissanke, Christoph Weniger, Andrew R. Williamson, and Gilles Louppe. Lightning-Fast Gravitational Wave Parameter Inference through Neural Amortization. In Third Workshop on Machine Learning and the Physical Sciences, 10 2020. arXiv: 12931
[68] Stephen R. Green and Jonathan Gair. Complete parameter inference for GW150914 14 using deep learning. Mach. Learn. Sci. Tech,, 2(3):03LT01, 2021. arXiv:2008. 03312 doi : 10. 1
[69] Maximilian Dax, Stephen R. Green, Jonathan Gair, Michael Deistler, Bernhard Schölkopf, and Jakob H. Macke. Group equivariant neural posterior estimation, International Conference on Learning Representations, 11 2022. :exiv:2111. 13139
[70] Chayan Chatterjee, Linging Wen, Damon Beveridge, Foivos Diakogiannis, and Kevin Vinsen.
Rapid localization of gravitational wave sources from compact binary coalescences using deep learning, arXiv preprint arXiv:2207.14522, 7 2022. DrXiv:2201 14522
14

---

[71] Stefan T. Radev, Ulf K. Mertens, Andreass Voss, Lynton Ardizzone, and Ullrich Köthe. Bayesflow: Learning complex stochastic models with invertible neural networks, 2020.
arXiv: 2003::20000000 06281
[72] Conor Durkan, Artur Bekasov, lain Murray, and George Papamakarios. Neural spline flows. In Advances in Neural Information Processing Systems, pages 7509-7520, 2019. arXiv
1906. 04032
[73] B.P. Abbott al. Observation of Gravitational Waves from Binary Black Hole Merger. Phys. Rev. Lett. 116(6)066112, 2016. arXiv: 03837
doi: :oi::00.103//10 100.1//0//yysse.ettt
116.061102
[74] Jonas Wildberger, Maximilian Dax, Stephen R. Green, Jonathan Gair, Michael Pürrer, Jakob H. Macke, Alessandra Buonanno, and Bernhard Schölkopf. Adapting to noise distribution shifts in flow-based gravitational wave inference. Phys. Rev. D, 107(8):084046, 2023. arXiv
2211.08801, doi: : 10. 1103/PhyssevD......m 107.084045
[75] Taco Cohen and Max Welling, Group equivariant convolutional networks. In Maria-Florina Balcan and Kilian Q. Weinberger, editors, Proceedings of the 33nd International Conference on Machine Learning, ICML 2016, New York City, NY, USA, June 19-24, 2016, volume 48 of JMLR Workshop and Conference Proceedings, pages 2990-29929 JMLR.org. 2016. URL:
http://prooceeding.. mlr press//448cohenc16 www.nc16..... w.w.ncl6.hhee html
[76] Mark Hannam, Patricia Schmidt, Alejandro Bohé, Leila Haegel, Sascha Husa, Frank Ohme, Geraint Pratten, and Michael Pürrer. Simple model of complete precessing black-hole- binary gravitational waveforms. Phys. Rev. Leut. 113:151101 101, Oct 2014. URL: https:// link. wwtps//rg////o/..o//////
/oi/10. www.////yyeeestttt.oom 113. 113.1511011 151101
doi : 10. 11100////aaeeeeetttt
113.151101
[77] Sebastian Khan, Sascha Husa, Mark Hannam, Frank Ohme, Michael Pürrer, Xisco Jiménez Forteza, and Alejandro Bohé. Frequency-domain gravitational waves from nonprecessing black-hole binaries. A phenomenological model for advanced detector era. Phys. Rev., D93(4):044007, 2016. arxiiv 07253
doi 10. 1ww.//////hhe......mm 93. 044007
[78] Alejandro Bohé, Mark Hannam, Sascha Husa, Frank Ohme. Michael Pürrer, and Patricia Schmidt. PhenomPv2 - technical notes for the LAL implementation. LIGO Technical Document, LIGO-TI500602-v v4, 2016. URL:
[79] Benjamin Farr, Evan Ochsner, Will M. Farr, and Richard O' Shaughnessy. A more effective coordinate system for parameter estimation of precessing compact binaries from gravitational waves. Phys. Rev. D, 90(2):024018, 2014. arXiv: 1404. 7070 doi :oi:10.. 100.10000///hh .90
024018
[80] J.R. Dormand and P.J. Prince. A family of embedded runge-kutta formulae. Jour- nal of Computational and Applied Mathematics, 6(1):19-26, 1980.
URL: hhtpps:// doi: Noi:https://
com wwt.://cc//////eree/eeeeeaae www www.ahannccciccc..ooooooooaooooa doi oal.:r//000..0000777700000000000
15

---

A
Gaussian flow
We here derive the form of vector field (0) that restricts the resulting continuous flow a one dimensional Gaussian with mean ft variance 0², With the optimal transport path (4 = t0ı,
(+( -1 mmin de from [6],, the sample-condttioal probability path (5) reads
Pe(0101) _
(10)
We set our target distribution
(1(01) = Via (()).
(11)
To derive the marginal probability path and the marginal vector field need two identities for the convolution of Gaussian densities. Recall that the convolution of two function is defined by * g(x) f(x yyg(y) dy. We define the function
ou.o² (0) = NNN [u, (0).
(12)
Then the following holds
Nip1, ofl * *N Nu, oद11 = NNI1 + 12 of + o
(13)
of
90 Nu2, o3]
(2210]+03 N222222 of )
(14)
o oz
Marginal probability paths
Marginalizing over 01 in 10) with find
/ pc (010))000))
(0(
/ N [10 of] (0) N [" (0 )d0
/ N [oo of (0 t0:) (tt (to)) d0
(15)
N [0 (0 0 M [t), (to) (09) d0,
N [4i of (16)²] (0)
where we defined 01 = 10, and used 13
Marginal vector field
now calculate the marginalized vector field (0) based on equation (8) in [66. Using the sample-conditional vector field (6) and the distributions (10) and we find
u4 (0) = Ue (0 01) (0000))((00)
(0) dB1
(01 (1 min )0
N t8,, of (0) M [n, (01) d0;
Pr(
(f( [ (0, (1 min
N [0. of (0 - 10:) / [Lū (ta) (101) d.
Øz
((( /
(0'
(1 O ) 0) N [0 (0 03) N [tp, (to)² (0) dB
0: t t
1
(-0" + (1 (1- mmin) (t) 0)
N | (0)) / [ (16) (0 0²) d0"
(e)
1.
1
/
(-0" + oe 0) N [0 (0 N [10, (16)² (0 0²) d0"
(()
T1 t
(16)
16

---

where we used the change of variables 8 = t01 and 8 = 89. Now we evaluate this expression using 12p. then the identities (33 and (4) and the marginal probability (15)
0
(0) pr(0) 01 1. (00.0? *N [ta [L,, (16)2) (0) p((0) (W of N [tū, (1t) (0)
-1 (0 lj) of
0
N [té, (o (tt) (0)
N tá. (o (tt)3) (0)
Pr(0)- ot + (to)²
p.(0)
(o) + (to)2) (0 t() 01 (0)
((( 1. (0) + (16)²)
(o? + (to)² 0,)0 / + tfi ot
t. (o - (t
(17)
By choosing a vector field 2y of the form 17 with learnable parameters P. 6², we can thus define continuous flow that restricted a one dimensional Gaussian.
B Mass covering properties of flows
In this supplement, we investigate the mass covering properties of continuous normalizing flows trained using mean squared error and in particular prove Theorem [] We first recall the notation from the main part. We always assume that the data is distributed according to p1 (6). In addition, there is known and simple base distribution Po and we assume that there is vector field u, : [0, x Rd - Rd that connects po and in the following sense. We denote by ór the flow generated by uz, i.e., dr satisfies
(e),(0) uo))
(18)
Then we assume that (61)-po = Pt and we also define the interpolations p, (dt)-po-
We do not have access to the ground truth distributions p, and the vector field u but we try to learn a vector field ey approximating ue- We denote its flow by 1/2 and define qe (v)) ( and qu Po- We are interested in the mass covering properties the learned approximation qı of p1. In particular, want to relate the KL-divergence DKL (PL((1llpp (l91) to the mean squared error,
10)={( -fo de f pd(())(u (,))) (
MSE, v) ==
(19)
of the generating vector fields. The first observation that without any regularity assumptions on 2 it is impossible to obtain any bound the KI KL-divergence in terms of the mean squared error.
Lemma 1. For every >> 0 there vector field u, and i2 and base distribution po go such that
MSE, (u, v) < and DRL(Pillg)) oc.
(20)
In addition can construct Un and vi such that the support of p1 is larger than the support of 91-
Proof. We consider the uniform distribution po 40 [[([-1,1]]
1)) and the vector fields
ue (0) 0
(21)
and
(0) fe for 0 S 0 <E otherwise.
(22)
As before, let di denote the flow of the vector field o, and similarly di denote the flow of e, Clearly (0) = 0. On the other hand
min(( et,e) if < <
(0) 0 otherwise.
(23)
17

---

In particular
if S 0 <,
(0)
(24)
0 otherwise.
This implies that p1 = (o ~ u[[-1,1] 1]). On the other hand 91 = (v/1) has support in [-1, 0J U [E, 1]. In particular, the distribution of qi is not mass covering with respect to p1 and (pu(p1 = co. Finally, we observe that the MSE can be arbitrarily small
(S v) = = [. dt (i((d)) u2( - 20(
f
1
E²
2
(25)
2
2
Here we used that the density of (d00) is 1/2 for -1 S 0 1
We see that an arbitrary small MSE-loss cannot ensure that the probability distribution mass covering and the KL-divergence is finite. On high level this can be explained by the fact that for vector fields V4 that are not Lipschitz continuous the flow is not necessarily continuous, and generate holes in the distribution. Note that chose Po to be uniform distribution for simplicity, but the result extends to any smooth distribution, in particular the result does not rely on the discontinuity of po-
Next, we investigate the mass covering property for Lipschitz continuous flows When the flows Ut and U, are Lipschitz continuous (in 0) this ensures that the flows 1h and Φ1 are continuous in and it is not possible to create holes in the distribution as shown above for non-continuous vector fields. We show weaker bound in this setting.
Lemma 2. For every 0 S 8 < there is a base distribution Pa qo and the are Lipschitt-comtinuuus
vector fields un and Ug such that MSE v) 8 and
(u.
(26)
Proof. We consider Po. 4o and ue as in Lemma and we define
20 for 0 0 <
(0)
2. for S 2E,
(27)
otherwise.
Then calculate for 0 s 0 S e -2 that
(0) = 2c
(28)
Similarly we obtain for ε < 0 5 28 (solving the ODE I' = 21)
1 (0) 2c (2s 0)e 24
(29)
We find
أ (0) = 0, (e =, (E) SE =2; v2(22) 2s.
(30)
Next we find for the densities of 91 that
1 se for 0 < 0 < e Φ1 (1/2 (0)) (o)))
(31)
2 e2 for 0 2s.
Together with (30p this implies that the density of gi is given by
-2 for < 0 < E,
(f) = 2 2 for 2s se < 0 < 2s.
(32)
Note that (0) 1/2 for -1 -1 0 < 1 and therefore
/ In (I(O)
1 Pi (d0) 1/¹ In(e" In(e2)
d0
(33)
qi (0)
2
18

---

and
2c
In Pi(O)
(d0)
f In(e
In(e=)) -86 -2
(34)
(1)
26-2e-2
22--0--2
Moreover note
91 (ds) = - (o(d))
1
25-50
2 -5(1-
Pı(dE),
(35)
2g
which implies (by positivity of the Xiveeeegge that
22-
(P(())
In
(1))) Pi (d0) 2 0.
(36)
We infer using also pı (0) = (() 1/2 for [-1, 0] n [2e, that
P1
(PL(lil)q)) In
Pi(0)) a(1 - -
(37)
(1)
On the other hand we can bound
1 / dt / J dt pc(d0) (r vn(0))
1² [u:(0)]2 /-- 3² ds
E3
(38)
We conclude that
DKL (1) (MSE, u)) u))
(39)
In particular, it is not possible to bound the KL-divergence by the MSE even when the vector fields are Lipschitz continuous.
Let us put this into context. It was already shown in 42 that we can, in general, not bound the forward -i--divereeece by the mean squared error and our Lemmas [ndd concrete examples. On the other hand, when considering SDEs the KI--divegencc can be bounded by the mean squared error of the drift terms as shown in [43] Indeed, in [42 the favorable smoothing effect was carefully investigated.
Here we show that we can alternatively obtain an upper bound on the KL-divergence when assuming that ue, th and Po satisfy additional regularity assumptions. This allows us to recover the mass covering property from bounds on the means squared error for sufficiently smooth vector fields. The scaling is nevertheless still weaker than for SDEs.
We state our assumptions. We denote the gradient with respect to 0 by V V and second derivatives by V
72 jur When applying the chain rule, we leave the indices implicit. We denote by
1/2
|- the Frobenius norm ||| Eig A3, of matrix. The Frobenius norm is submultiplicative i..,, |AB| < [A| [B] and directly generalizes higher order tensors.
Assumption 1. We assume that
Vud S L, Vol S L, /Vuu S L', V L.
(40)
We require further assumption on Po- Assumption 2. There constant C1 such that
(nnpo(() S (1 + (0).
(41)
We also assume that
Epo |0 < C2 < oo.
(42)
Note that (41) holds, e.g., if po follows a Gaussian distribution but also for smooth distribution with slower decay at co. If we assume that V In (n()) is bounded the proof below simplifies slightly.
This is, e.g., the case if (0)
[6] [0]
We need some additional notation. It is convenient introduce oi = th (#.)-1, i.e., the flow from time to & (in particular of Φx) and similarly for th We can now restate and prove Theorem
19

---

Theorem 2. Let Po and assume and vt are two vector fields whose flows satisfy p1 = =61))=P
and q1 (V1).90- Assume that po satisfies Assumption Lnd u, and v, satisfy Assumption 7 Then there is constart C' >> O depending on L, L/, C1, C2, and d such that (for MSF (u, v) < T))
DKL(Pill4) & C (S, v)a.
(43)
Remark We do not claim that our results are optimal, might be possible to find similar bounds for the forward KL-divergence with weaker assumptions. However, we emphasize that Lemma 2 shows that the result of the theorem not true without the assumption on the second derivative of 2 and Ue.
Proof. We want to control DKL (pill))) It can be shown that (see equation above (25) in 43 or Lemma 2.19 in [42]
0, (Pill,() = - - (cd - (x)
V In (0))
(44)
Using
can bound this by
< (f (:((d)))u. (000)))(_())00((00 )2((0)))(()))((0))))))))))))))
& (ʃ ((.(a0)))Iinxx(00 (0:000-VInnaa { (45)
a (P:lle))
We use the relation (see 3
In(p: (d)(()) = In(po((0)) - S.¹ (div (o))
(46)
which can be equivalently rewritten (setting = x((0)) as
(0)) = im(po(o)0)) - 1/" (div u. (o)
(47)
We use the following relation for V
(o' exp (L. dr (Vu)))6(00))
(48)
This relation is standard and can be directly deduced from the following ODE for Vol
0, ()) v0,6 (0) ((u((d(())) ((Vu))dd;(0)) (o' We can conclude that for 0 - s,t S 1 the bound
(49)
[oo s c
(50)
holds. We find
[in(p((0)) = v VIn(p)) (o - / div u.)()))) vol
(51)
[iI(p))( + L'et,
and similar bound holds for q:. In words, we have shown that the score of p: at @ can be bounded by the score of Po of theta transported along the vector field u, minus a correction which quantifies the change of score along the path. We now bound using the definition pt- = (ot)..P and the assumption
/
pr pr(d0) (
V (mmp@cch000)( = Ep |inpo(0o)1
(52)
< Ep
< 20 + Ex 10 1) < 2C + C?).
Similarly we obtain using qo Po
f pc((d) V In po((0))\V In [n
(53)
20

---

In words, to control the score of integrated with respect to we need to control the distortion we obtain when moving forward with u and backwards with v. We investigate (b)). We now show
O th 0ith (0) = (.(0) =u,(t) uv)
(54)
First, by definition of Φ. we find
On (0) 0 (0) = (e(bi0, =
= (O)
(55)
To evaluate the second contribution we observe
0 = 0,00 = 0n wnwit (0)1 - ( (0)1
(O )/1 (0) (0)| (O.
= (e + (1)) (56)
(0) + On 2 + h=
Now 54
follows from 555 and
56
together with of = v Id. Using 54 we find
(() (1/6 +
oi)n)) (Vv))(dr((o)) ((we - (0)))) (57)
Using 50 we conclude that
(B.) Bol = 1.' a. (e08sos ds $ { ('VV6))(0s - luu - vol())((u)) ds
(58)
S e S." lu= - v. (fo)) ds.
We use this and the assumption
to continue to estimate 53 as follows
J In (n - Vi
s <c po(d)) (1 (166.
$ <0 / po(d0a) (60) Pul 10,1)2
$ 33G + 3C3 / pP((dd) (10b): (Bo) - -bol + 10b1
s 30 + Epp 1013) 3C7e2L / po(dlb) (Lr ds lu, Ual v=l(
(59)
Here we used (a + b c)² < 3(a² + 62 + (²) in the second to last step. We bound the remaining integral using Sauchy-Schware as follows
/ po(d)0) ( [uo c=((6s(0h))
S po(c60) (L, ds lu, v.l (oa
(% ds 12
< %:- ds / po((66)lue - (8)))
l ds / poldd=)lus-uu - val((0.)
< / ds / (p.(d.))uu (0.) = MSE (u, v).
(60)
The last displays together imply
/ [ 303 (1 + Eno 10o1 +e 2h MSE,(u, v))
(61)
21

---

Now we have all the necessary ingredients to bound the derivative of the KL-divergence. We control the second integral in (5)) using 51 (and again (Ei=1 a;)² I 4 - ag) as follows,
/ Pi(d))\V (l)
(62)
22. e" + 4c f po(dt) (IV (o(o0)))) + Vinp()))0))))
Using 52 and 61 we finally obtain
(c(0))VInpp00 V In (0) 88 1.2 + CREEL (20(1 + C2) + 12 (u,
(63)
I C(1 MSE,(u,))
for some constant C' > O. Finally, we obtain
(ptl(pil19)== = 1/' dt (P..l
- (C(1+ MSE, u))) S (frido))uu0)) - u) 1
(64)
S (C(1 + MSE, o))} (fo dt I (1(0))))(00)((100))))) un 1
< (C(1 + (M, 0))} (u,
C SBI Benchmark
In this section, we collect missing details and additional results for the analysis of the SBI benchmark in Section
C.1 Network architecture and hyperparameters
For each task and simulation budget in the benchmark, we perform mild hyperparameter optimiza- tion. We sweep over the batch size and learning rate (which is particularly important as the simulation budgets differ by orders of magnitudes), the network size and the parameter for the time prior defined in Section 3.3 (see Tab. 2 for the specific values), We reserve 5% of the simulation budget for validation and choose the model with the best validation loss across all configurations.
C.2 Additional results
We here provide various additional results for the SBI benchmark. First, we compare the performance of FMPE and NPE when using the Maximum Mean Discrepancy metric (MMD). The results can be found in Fig. 6 FMPE shows superior performance to NPE for most tasks and simulation budgets. Compared to the C2ST scores in Fig. the improvement shown by FMPE in MMD is more substantial.
Die-Zcompanc the FMPE results with the optimal transport path from the main text with comparable score matching model using the Variance Preserving diffusion path [[5]. The score matching results were obtained using the same batch size, network size and learning rate the FMPE network, while optimizing for 8min {0.1, 1, 4] and Bmax E (4, 7, 10). FMPE with the optimal transport path clearly outperforms the score-based model on almost all configurations.
In Fig. we compare FMPE using the architecture proposed in Section 3.2 with (t. condiitiio via gated linear units to FMPE with a naive architecture operating directly on the concatenated (t, 0,,r) .r) vector. For the two displayed tasks the context dimension (im(x) = 100 much larger than the parameter dimension dim(@) E {5, 10). and there is a clear performance gain in using the
GLU conditioning. Our interpretation is that the low dimensionality of (1, 0) means that is not well-learned by the network when simply concatenated with z.
22

---

hyperparameter
sweep values
hidden dimensions number of blocks
2 for n (4....,10) 10}
10..... 18
batch size learning rate a (for time prior)
2 for E {2..... .,, 9}
L.-3, 5.e-4, 2..-4, 1.e4
-0.25, -0.5, 0, 1,
Table 2: Sweep values for the hyperparamters for the SBI benchmark. We split the configurations according to simulation budgets, e-g. for 1000 simulations, we only swept over smaller values for network size and batch size. The network architecture has a diamond shape, with increasing layer width from smallest to largest and then decreasing to the output dimension. Each block consists of two fully-connected residual layers.
0.20
GL
GL-U
GM
Two Moons
SLCP
NPE
0.10
EMPE
MMD
0.00
B-GLM
G------ww
SLCP-D
SIR
0.4
LV
0.2
0.0
10
10
10a
10a
10a
Number Simulations
Figure 6: Comparison of FMPE and NPE performance across 10 SBI benchmarking tasks 46 here quantify the deviation in terms of the Maximum Mean Discrepancy (MMD) an alternative metric to the C2ST score used in Fig. 4 MMD can be sensitive its hyperparameters [46]. so we use the C2ST score primary performance metric.
We
Fig. 9 displays the densities of the reference samples under the FMPE model as histogram for all tasks (extended version of Fig. 3). The support of the learned model q(0 ()) covers the reference samples 0 p(0]z), providing additional empirical evidence for the mass-covering behavior theoreti- cally explored in Thm. However, samples from the true posterior distribution may have small density under the learned model, especially if the deviation between model and reference is high; see Lotka- Volterra (bottom right panel). Fig. displas P-P plots for two selected tasks.
Finally, we study the impact of our time prior re-weighting for one example task in Fig. II We clearly see that our proposed re-weighting leads to increased performance by up-weighting samples for closer during training.
23

---

1.0
1.0
1.0
GL-U
GM
Two Moons
0.8
SMPE
0.8
0.8
SLCP
FMPE
0.6
0.6
0.6
C2ST
1.0
1.0
1.0
1.0
1.0
B-GLM
B-GLM-Raa Raw
SIR
0.8
0.8
0.8
SLCP-D
0.8
0.8
LV
0.6
0.6
0.6
0.6
0.6
10
10
10
10a
10
10
10
10
10
10a
10
10
10a
10*
10
Number of Simulations
Figure 7: Comparison of FMPE with the optimal transport path (as used throughout the main paper) with comparable models trained with a Variance Preserving diffusion path [[5] by regressing on the score (SMPE). Note that the SMPE baseline shown here is not directly comparable to NPSE (8 [9] as this method uses Langevin steps, which reduces the dependence of the results on the vector field for small (at the cost of tractable density).
1.0
B-GLM-Raw
0.9
embeddiing
Concatenation
SLCP-D
0.7
0.6
0.5
10
10
10
10a
10
105
Number of Simulations
Comparison of the architecture proposed in Section 3.2 with gated linear units for the (L, 0) condiiiiig (red) and naive architecture based simple concatenation of (1,0, (black). FMPE with the proposed architecture performs substantially better.
Figure 8:
24

---

GL
(~p()(a)
(~0((0))
-2
GL-U
15
10
-5
-15
10
-5
15
-10
5
GM
-3.0
2.5
0.0
2.5
5.0
5
0
Two Moons
-5.0
-2.5
SL.CP
9
8
7
10.0
-7.5
-5.0
-25
10
55
B-GLM
20
0
-10
-10
B-GLM-REE
15
10
-5
10
-5
SLCP-D
100
8
-6
10.0
7.5
-5.0
-25
- 10
-5
0
SIR
-7.5
-5.0
2.5
0.0
-6
2
-4
-2
2
LV
-10
-5
-10
5
-3.0
2.5
2.5
log q(((z))
Figure 9: Histogram FMPE densities q@g xx) for samples 0 ~ q(((2)) and reference samples 0 p(0[zr) for simulation budgets N 10 (left), N = 10 (center) and N 10 (right). The reference samples ~ p(0|z)) are within the support of the learned model g(0)=)), indicating mass covering FMPE results. Nonetheless, reference samples may have small density under x). if the validation loss high, see Lotka- Voltera (L.V).
25

---

1.0
1.0
1.0
CDF(p)
CDF(p)
CDF(p)
0.5
0.5
0.5
0.0
0.0
0.0
1
0
1
p
P
P
1.0
1.0
1.0
CDF
CDF(p)
CDF(p)
0.5
0.5
0.5
0.0
0.0
0.0
0
p
P
p
Figure 10: P-P plot for the marginals of the EMPE-posterior for the Two Moons (upper) and SL.CP (lower) tasks for training budgets of 103 (left), 10 (center), and 10 (right) samples.
Two Moons
0.9
law umiform
0.8
0.7
0.5
10
10"
10
Number of Simulations
Figure 11: Comparison of the time prior re- re-weighting proposed in Section 3.3 with uniform prior over on the Two Moons task (Section 4. The network trained with the re-weighted prior clearly outperforms the reference on all simulation budgets.
26

---

D Gravitational-wave inference
We here provide the missing details and additional results for the gravitational wave inference problem analyzed in Section 5
D.1 Network architecture and hyperparameters
Compared to NPE with normalizing flows, FMPE allows for generally simpler architectures, since the output the network simply vector field. This also holds for NPSE (model also defined by a vector) and NRE (defined by scalar). Our FMPE architecture builds on the embedding network developed in 77. however we extend the network capacity by adding more residual blocks (Tab, 3 top panel), For the (t, -)-codiitigning we use gated linear units applied to each residual block, as described in Section 3.2 We also use a small residual network to embed (t, 0) before applying the gated linear units.
In this Appendix we also perform an ablation study, using the same embedding network as the NPE network (Tab. [] bottom panel). For this configuration, we additionally study the effect of conditioning on (1, () starting from different layers of the main residual network.
D.2 Data settings
We use the data settings described in 7 with few minor modifications. In particular, we use the waveform model IMRPhenomPv2 76 and the prior displayed in Tab. 4 Generation of the training dataset with 5,000,000 samples takes around 1 hour 64 CPUs. Compared to 77, reduce the frequency range from [20, 1024 Hz to [20, 512] Hz to reduce the computational load for data preprocessing. We also omit the conditioning on the detector noise power spectral density (PSD) introduced in (7 as we evaluate on a single GW event. Preliminary tests show that the performance with PSD conditioning is similar to the results reported in this paper. All changes to the data settings have been applied to FMPE and the NPE baselines alike to enable a fair comparison.
D.3 Additional results
Tab. displays the inference times for FMPE and NPE. NPE requires only single network pass to produce samples and (log- probabbpbbbb,,sssss,, whereas many forwards passes are needed for FMPE to solve the ODE with specific level of accuracy. A significant portion of the additional time required for calculating (log-)probabilities in conjunction with the samples is spent on computing the divergence of the vector field, Eq. (3).
hyperparameter
values
residual blocks
2048, 4096 x 3, 2048 x 3, 1024 x 6, 512 x 8, 256 x 10, 128 5, 64 x 3, 32 3, 16 x6xx3
residual blocks (t, embedding batch size
16, 32, 64, 128, 256 4096
learning rate
5.c-4
a (for time prior)
residual blocks
2048 x 2, 1024 x 4, 512 x 4, 256 x 4, 128 x 4, 64 x 32 x 3, 16 3
residual blocks (t, embedding
16. 32, 64, 128, 256 4096 5.e-4
batch size learning rate a (for time prior)
Table 3: Hyperparameters for the FMPE models used in the main text (top) and in the ablation study (bottom, see Fig. (2). The network is composed of sequence of residual blocks, each consisting of two fully-connected hidden layers, with a linear layer between each pair of blocks. The ablation network is the same as the embedding network that feeds into the NPE normalizing flow.
27

---

Description
Parameter
Prior
component masses
mi. m2
[10, 120] Mo. mi 2 m2
chirp mass ratio
Mc = (m1m²) (m1 + m2)
[20. 120) Mg (constraint)
q m2/mi
[0.125, 1.0] (constraint)
spin magnitudes spin angles time of coalescence
42
[0, 0.99
11, 02. 012, DJI.
standard as in (79
[-0.03, 0.03] 8
luminosity distance
dp
1000, 1000] Mpc
reference phase
d.
0, 2
inclination polarization
JN
00, uniform in sine
/b
0.
sky position
o,,
uniform over sky
Table 4: Priors for the astrophysical binary black hole parameters. Priors are uniform over the specified range unless indicated otherwise. Our models infer the mass parameters in the basis (Mer 4) and marginalize over the phase parameter de.
Fig. 12 presents comparison of the FMPE performance using networks of the same hidden dimen- sions as the NPE embedding network (Tab. bottom panel). This comparison includes an ablation study the timing of the (t. 8) GLU-conditioning. In the top-row network, the (t, 0) conditioning is applied only after the 256-dimensional blocks. In contrast, the middle-row network receives (t, 0) immediately after the initial residual block. With FMPE we achieve performance comparable NPE, while having only ~ 1/3 of the network size (most of the NPE network parameters are the flow). This suggests that parameterizing the target distribution in terms of vector field requires less learning capacity, compared to directly learning its density. Delaying the (t. 0) conditioning until the final layers impairs performance. However, the number of FLOPs inference is considerably reduced, as the context embedding can be cached and a network pass only involves the few layers with the (t, 0) conditioning. Consequently, there's trade-off between accuracy and inference speed, which we will explore in greater scope in future work.
Network Passes
Inference Time (per batch)
FMPE (sample only)
248
268
FMPE (sample and log probs)
350
352s
NPE (sample and log probs)
1.5s
Table 5: Inference times per batch for FMPE and NPE on a single Nvidia A100 GPU, using the training batch size of 4096. We solve the ODE for FMPE using the dopri5 discretization [80] with absolute and relative tolerances le-7, For FMPE, generation the (log-)probabilities additionally requires the computation of the divergence, therefore limits the maximum batch size that can be used at inference.
equation (3). This needs additional memory and
mn
111
9
4
1/2
4
4
d
20
FMPE late GLU
45.3
41.5
1.2
0.5
18.0
11.8
0.3
12.3
8.1
11.4
7.5
0.4
(SD)
FMPE carly GLU
0.8
1.0
0.2
0.3
1.7
13.0
0.4
10
mmat
NPE
1.2
2.5
3.2
1.6
0.8
0.4
0.3
4.4
9.1
10.1
8.6
0.6
Figure 12: Jensen-Shannon divergence between inferred posteriors and the reference posteriors for GW150914 (73). We compare two FMPE models with the same architecture the NPE embedding network, see Tab. bottom panel. For the model in the first row, the GLU conditioning (8, t) is only applied before the final 128-dim blocks. The model in the middle row is given the context after the very first 2048 block.
28


Reviews:
[{"review_id": "mdGCNFioVt", "reviewer": "Reviewer_1iUK", "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. \n\nThe 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. \n\nFurthermore, 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. \n\nIf 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.\n\nIn 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.", "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).\n- 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\".\n- It was not super clear to me why you have to sample $\\theta_1$ and what does this mean in the whole setup\n- Can you give more intuition in to the failure cases for the assumptions in Theorem 1?\n- 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.\n- 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?\n- 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?\n- Line 313: What do you mean with the much simpler vector field of FMPE? Can you give more elements into this comment?", "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. ", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- Very good writing and style of explanations.\n\n- Simple (in the sense of not being very far reached from previous works) yet relevant new method for simulation-based inference.\n\n- 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?\n\n- Some mathematical notation are slightly inconsistent and hard to follow (I will be more precise in my questions)\n\n- Saying that the method is **scalable** without being very precise in which sense -- in observation space, not parameter space.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the thorough report and are pleased to hear that they enjoyed reading the paper.\n\n### Weaknesses\n\n* *[“Results on real data seem rather preliminary”, \"Why [...] use FMPE for it then?\"]* Indeed, results are not comprehensive from a physics perspective, and we plan to carry out a more in-depth analysis in the future (l.318-321). Nevertheless, FMPE is highly promising as the first generic method to achieve the desired accuracy for GW inference (i.e., not using group equivariances) and unlike GNPE it allows for density estimation. We believe that it is reasonable to expect that once symmetries are incorporated into FMPE, it should surpass GNPE performance. This extension is planned for a follow-up study. We will update the discussion Sec. 5.3 to clarify this.\n\n* We will ensure consistency of notation in the updated version.\n\n* We indeed intended to refer to scalability in observation space, and will clarify this in the updated paper. [Note that we additionally expect FMPE to scale to high dimensional parameter spaces, as this property has been extensively demonstrated in the context of generative modeling. However, we did not show this in this paper.]\n\n### Questions\n\n* *[Training time vs simulation budget.]* This depends on the problem at hand. For gravitational waves, our training time ($\\approx$2 days on 1 GPU, 32 CPUs) is substantially larger than the simulation time ($\\approx$1 hour on 64 CPUs), making the former the bottleneck. Simulation efficiency is however an important consideration in other SBI problems. We will include the simulation time in App. D.2.\n\n* *[Parameter space vs sample space.]* We agree that in the SBI context \"parameter space\" is the common terminology for $\\theta$. We used \"sample space\" to be consistent with the flow matching/diffusion literature, and we will clarify this in a footnote.\n\n* *[Sampling $\\theta_1$.]* $\\theta_1$ refers to a parameter sample from the training dataset, sampled from the joint distribution $(\\theta_1, x) \\sim p(\\theta_1, x)$ in the standard SBI fashion. $\\theta_1$ is then used to condition the transport paths between base and target distribution, as proposed by [1]. This way, we effectively bypass the need to estimate the otherwise intractable marginal vector field between the base and the target distributions.\n\n* *[Theorem 1 failure cases.]* The assumptions fail to hold, e.g., when non-differentiable activation functions are used in the architecture.\n\n* *[Impact of choice of loss, architecture, and training procedure on scalability.]* Thank you for pointing out that we are missing this; we will add a paragraph in the discussion. Briefly, scalability is improved by regressing on a vector field rather than an entire distribution: we empirically find that fewer network parameters are needed to learn this vector field. Given the same computational budget, this allows for larger networks, ultimately enabling them to solve more complex problems (see also final response below). Furthermore, our architecture for FMPE (a straightforward ResNet with GLU conditioning) facilitates parallelization and allows for cheap forward/backward passes.\n\n* *[NPE setup for benchmark? Poor FMPE Lotka-Volterra performance?]* We did not run the NPE baseline tests ourselves, but reported the results from [2]. As pointed out in [2], no existing SBI method successfully solves Lotka-Volterra (Sec. 3, #2). Ref. [2] does not comment on the reasons for this. The Lotka-Volterra model is defined by a set of coupled differential equations (see T.10 in [2]), and presumably the pattern between model parameters and observations induced by these is hard to learn.\n\n* *[Why does FMPE train faster than other methods?]* The key difference is that FMPE parametrizes a vector field, while NPE parametrizes a complex distribution (see Fig. 1). When using neural spline flows [3] for NPE, a likelihood evaluation requires application of the coupling transforms (Sec. 2.1 in [3]), which involves the evaluation of splines for each flow layer. The flow layers are also repeatedly conditioned on the context data. FMPE in contrast simply regresses on a vector field during training, which amounts to a simple neural network pass, taking the context data as input, which is computationally faster. We will clarify this in the updated paper.\n\n* *[What is meant by \"simple\" vector field on l.313?]* Thank you for this question, as it forced us to think more precisely about what we mean. In fact, there are three main aspects that make the continuous NF (CNF) architecture used in FM simpler than the NF: (1) In each transform step the CNF applies a translation, whereas the NF uses an elementwise spline; (2) Although both networks use (in our case) ResNets, the NF requires a separate ResNet for each step, whereas a single `t`-dependent ResNet suffices for the CNF; and (3) The connection between network outputs and the transform is more transparent for the CNF. Although all of these contribute to \"simplicity\", we believe that remark (2) is key to the reduction in required (computational) complexity for good results. We will expand upon the discussion in the paper to clarify.\n\n### Limitations\n\nWe will add a short section on limitations just before the conclusions. We have now performed calibration tests (P-P plots) for the SBI benchmark task and we include representative examples in the PDF (see main rebuttal). Those results were obtained as described in  [4].\n\n### References\n\n[1] Lipman, Yaron, et al. \"Flow matching for generative modeling.\" arXiv preprint arXiv:2210.02747 (2022).\n\n[2] Lueckmann, Jan-Matthis, et al. \"Benchmarking simulation-based inference.\" International conference on artificial intelligence and statistics. PMLR, 2021.\n\n[3] Durkan, Conor, et al. \"Neural spline flows.\" Advances in neural information processing systems 32 (2019).\n\n[4] Talts, Sean, et al. “Validating Bayesian Inference Algorithms with Simulation-Based Calibration.” arXiv preprint arXiv:1804.06788 (2018)\n"}, {"review_id": "qxCdwwU4uZ", "reviewer": "Reviewer_SM25", "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.\n\nSecondly, 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.\n\nIn 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.\n\nThe 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.\n\nLastly, the paper presents several experiments designed to demonstrate the effectiveness of the proposed method, including gravitational-wave inference tasks.", "questions": "N/A", "limitations": "Please refer to the comments provided in the weaknesses section.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "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.\n\nThe 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\nWe thank the reviewer for the careful reading of our manuscript and positive comments on our writing and perspective. We fully agree with the reviewer that the main challenge was not purely technical; rather, **we combine existing methods in a new way, and show that this can solve a very hard task for which other generic SBI methods fail.** We believe that a major advantage of our approach lies in its straightforwardness, making it more accessible and easier to apply.\n\n* *[Provide more detail on technical challenges.]* Although flow matching for SBI is a natural extension, there are several caveats in applying a technique originally proposed for generative modeling. A central contribution of our paper is to point these out, analyze them in detail (theoretically and empirically), and provide improved design suggestions where necessary. This includes for example:\n\n    1. **Dimensionality of sample space.** Traditional successes of flow matching (and diffusion models) have been powered by architectures tailored for high-dimensional sample spaces. In SBI, by contrast, the observation space (or context) is typically high-dimensional, whereas the parameter space (or \\(\\theta\\)-space) is low-dimensional. This requires a different type of architecture (Sec. 3.2). \n\n    2. **Focus on distribution vs. individual samples.** Flow matching has been shown to excel at generating high-fidelity individual samples. However, in SBI, our primary concern is to accurately capture the entire distributional shape. We demonstrate that flow matching also excels at this task.\n\n    3. **Conservative in estimates**: SBI requires conservative estimates to avoid dismissing scientifically plausible scenarios. We explored this aspect in detail from both theoretical and empirical standpoints, showing that FMPE does indeed provide probability mass covering estimates.\n\n* *[Gravitational-wave inference task demanding enough?]* GW inference is in fact an *extremely* challenging task (incl. high-dimensional observations, complicated physics models, low signal-to-noise ratio), with very high accuracy requirements (see l.301-302) and the potential for great scientific impact (also see the general reply). The only SBI method shown to meet these requirements relies on a symmetry-aware framework and thereby solves an arguably easier problem (l.276-280). FMPE is the only *generic* SBI method with comparable accuracy, and substantially outperforms other competing methods (see general reply). \n\n* *[More demanding example as platform for discussing technical contributions?]* As mentioned above, GW inference is a great platform to push the limits of SBI algorithms, and we discuss how our proposed network architecture contributes to the superior performance (l.310-317). We will expand this discussion in the updated version of the paper. \n"}, {"review_id": "X8vXsRwVBA", "reviewer": "Reviewer_7UFh", "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.\n\nUnlike ELBO-based training methods, the proposed approach does not ever sample from the model during training.\n\nExperiments on some benchmark problems seem to be done, though this reviewer has no experience in these problems.", "questions": "  - How is JSD computed? This would typically require sampling from the true posterior. \n\n  - Are the true posteriors known for these problems?\n\n  - 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?\n\n  - 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?\n\n  - In regards to Theorem 1, in the author's opinions, what is the most problematic assumption that may not hold in practice?", "limitations": "I do not see a discussion on limitations.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "  - Well-written, motivated, and easy to understand.\n\n  - Method is straightforward but seems to work well.\n\n  - Theorem 1 is great.", "weaknesses": "  - Benchmarks do not seem to be be challenging.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive review and helpful comments. We are happy that they liked our Theorem 1!\n\n### Weaknesses\n\n* *[Benchmarks not challenging.]* As we argue in the response to all reviewers, the gravitational wave (GW) inference problem (Sec. 5) is in fact extremely challenging. We show that FMPE is the first generic (i.e., non symmetry-enhanced) SBI method that meets the accuracy requirements, significantly outperforming NPE. Regarding the SBI benchmark test suite (Sec. 4), we agree that this contains some simpler toy problems, but this is intended more to provide a broad set of tests relevant to SBI algorithms.\n\n### Questions\n\n* *[How is JSD computed?]* In all cases we have access to reference posteriors: the SBI benchmark provides these in terms of MCMC samples, and for the GW problem we obtain samples using the method described in [1], which uses the exact likelihood corresponding to stationary Gaussian noise. \n \n* *[Inference problems challenging enough?]* See reply to weaknesses above.\n\n* *[Fitting an approximate posterior via ELBO?]* There are many ways to fit an approximate posterior. As you point out, one could use variational methods, which minimize KL(p_model || p_true_posterior). Alternatively, one could use NPE (which we provide as a baseline), which optimizes KL(p_true_posterior || p_model). Finally, one could use FMPE and optimize Eq. 7. All three of these methods directly fit the posterior (and will asymptotically recover the exact posterior, given enough training data and network capacity). **FMPE outperforms both of these alternatives, NPE (see Sec. 5) and ELBO maximization (see [2]).** Ref. [2] applies variational inference to the GW example, and FMPE achieves considerably more accurate results (see e.g., Fig. 2 in [2], and note that they only present results on simulations, but not on real data).  \n\n* *[Sampling from the prior?]* The GW problem is indeed too challenging to be solved by sampling from the prior. In our pipeline we only ever sample from the prior to generate the training dataset. (This enables amortized inference across the entire prior.) The final model is conditional on the observational data and thereby estimates the posterior, not the prior (Eq. 7). \n\n* *[Bayesian neural networks?]* This is an interesting suggestion, and we would be curious to see whether flow matching can be applied to Bayesian neural networks. In this work, however, we specifically investigate SBI, which has very different requirements than fitting Bayesian networks. Our results therefore provide no indication on the feasibility of this.\n\n* *[Constant in Theorem 1]* Yes, the constant grows exponentially with time. This is the expected behavior (see, e.g., Groenwall’s inequality). In our setting we only consider the time interval [0,1]. We have not tried to estimate the constant for the experimental problems, which is a difficult task because this requires estimates of the derivatives of the unknown true posteriors.   \n\n* *[What assumptions in Theorem 1 may not hold in practice?]* When parametrizing the learned vector field with a neural network the derivative and the second derivative will typically only be bounded by large constants $L$ and $L’$ (or unbounded when using non-differentiable activation functions like Relu). Then the bound will be weak. Also, it is difficult to find $L$ and $L’$ in practice. \n\n### References\n\n[1] Dax, Maximilian, et al. \"Neural Importance Sampling for Rapid and Reliable Gravitational-Wave Inference.\" Physical Review Letters 130.17 (2023): 171403.\n\n[2] Gabbard, Hunter, et al. \"Bayesian parameter estimation using conditional variational autoencoders for gravitational-wave astronomy.\" Nature Physics 18.1 (2022): 112-117.\n"}, {"review_id": "gnUUTEAAuE", "reviewer": "Reviewer_VmG3", "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.\nThe 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.\nThe 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.", "questions": "- Can we have a reference for the computation of Eq (54) in App B ? (just like (48) was derived from (49))\n- As mentioned in Sec 3.1, [1] provides a MSE bound of the KL in the SDE framework. Did you try such models ?\n- Can you further illustrate the superiority of your time prior ?\n- 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 ?\n- How many integration steps were used in the experiments when sampling or computing the likelihood ? Was the dopri5 integrator always used ?\n- 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 ?)\n- The training time difference between NPE and FMPE in GW (L294-L296) is very surprising. How do you explain it ?\n- What is the design of the Normalizing Flows for NPE experiments ?\n\n[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.", "limitations": "Not applicable.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- Well written paper\n- 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.\n- 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).\n- The experimental results (see Fig 3 in Sec 4 for instance) don't convey a clear superiority of the proposed method over NPE.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive report and helpful comments.\n\n### Weaknesses\n\n* *[Cost of likelihood evaluations.]* We agree that flow matching likelihood evaluations are expensive due to the many backward network passes required (taking ~6 minutes for the gravitational waves example). However, we disagree that this is generally prohibitive for SBI since likelihood evaluations are not needed for training when using the sample-conditional flow matching objective (4) (see also (7)). Moreover, for applications such as GW analysis, training and data simulation tasks are often much more costly than inference (training ~ days; inference with likelihoods ~ hours). We agree that the paper will benefit from an expanded discussion of this, so we will add a short paragraph in a new limitations section before the conclusions.\n\n* *[SBI benchmark density estimation.]* We found the histograms useful for illustrating when and how FMPE inaccuracies manifest, as well as to enable verification of mass coverage. Fig. 4, in particular, shows a failure case for the Two Moons task due to a small simulation budget $N = 10^3$, which goes away for larger $N$. For the majority of tasks, accuracy improves significantly with larger $N$, while for Lotka-Volterra, performance is always poor (as it is for NPE as well). The performance shown in the histograms is consistent with the c2st scores of Fig. 3, which we claim in general is competitive with NPE. We will add a statement to Sec. 4 on the consistency between c2st and histograms.\n\n* *[FMPE vs NPE for SBI benchmark.]* Because of the simplicity of the benchmark tasks, they are generally used by the SBI community as a sanity check rather than a discriminative test. Here we are not making a strong claim about superiority of FMPE vs NPE on the benchmark (l.232--237: \"FMPE exhibits comparable performance to an NPE baseline model for most tasks and outperforms on several\"). As pointed out in [1], similar statements are true for *any* SBI method (e.g., Sec. 3, #5): “no single method (NLE, NRE, NPE) outperforms another method on all tasks.” A central claim of our paper is the scalability to complicated observation spaces, which is depicted by the GW example (see Fig. 5, which demonstrates clear superiority over NPE).\n\n### Questions\n\n* We will add a more detailed derivation of Eq. (54). The key observation is that \n$\\partial_h \\psi_t (\\psi_{t+h})^{-1}(\\theta)\\rvert_{h=0} = - v_t(\\theta) $  which can be obtained by applying $ \\partial_h $ to $ \\theta = \\psi_{t+h} (\\psi_{t+h})^{-1}(\\theta) $ at $ h=0 $.\n\n* *[SDE models.]* We here focus on continuous flows with likelihood access, which the SDE formulation does not provide. As the GW example requires likelihoods for importance sampling, we have not applied SDE frameworks.\n\n* *[Superiority of new time prior.]* We have included in the PDF (attached to main rebuttal) an ablation of our time prior on the Two Moon tasks, which shows a significant performance improvement for the power-law prior (Sec 3.3).\n\n* *[SBI benchmark histograms.]* See above for detailed comments on Figs. 4 and 9. We used the NPE results reported in [1] in the paper, but we have now re-run these analyses ourselves to include the Two Moons NPE histogram in the PDF. The NPE histogram in this case looks better than the FMPE histogram, which is consistent with the c2st score for $N = 10^3$ (Fig. 3).\n\n* *[ODE integration.]* We always use the dopri5 solver to generate samples and likelihoods. The number of integration steps for the GW inference task is given in Table 5 in App. D (sample only: 248; sample and log probabilities: 350).\n\n* *[FM likelihood costs.]* For the results in the initial submission we computed the divergence exactly. We have now experimented also with Hutchinson’s trace estimator, and found that given a fixed number of integration steps, it provides a speed-up by a factor of $\\approx$4. However, in practice, one may then need to use more integration steps to keep the error sufficiently small for downstream applications. We further note that improving the sampling times of diffusion models/flow matching is an active area of research and promising approaches have been proposed for further speed-ups [2,3].\n\n* *[Training time improvements.]* The key difference is that FMPE parametrizes a vector field, whereas NPE parametrizes a complex distribution via a parametrized mapping on the sample space (Fig. 1). When using neural spline flows [4] for NPE, a likelihood evaluation requires application of the coupling transforms (Sec. 2.1 in [4]), which involves the evaluation of splines for each flow layer as well as repeated conditioning on the context data. FMPE in contrast simply regresses on a vector field during training, which amounts to a simple neural network pass, which in turn is computationally faster. We will clarify this in the updated paper.\n\n* *[NPE architecture]* All NPE results reported in this paper use neural spline flows [4]. For the benchmark (Sec. 4), we report the results from [1] (see Sec. H.5 for their discussion on hyperparameters). For the GW example (Sec. 5), we ran NPE ourselves, using the configuration in [5] (see Sec. Supplemental Material => Training Data => Neural Network for their details).\n\n### References\n\n[1] Lueckmann, Jan-Matthis, et al. \"Benchmarking simulation-based inference.\" International conference on artificial intelligence and statistics. PMLR, 2021.\n\n\n[2] Lu, Cheng, et al. \"Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps.\" Advances in Neural Information Processing Systems 35 (2022): 5775-5787.\n\n[3] Shih, Andy, et al. \"Parallel Sampling of Diffusion Models.\" arXiv preprint arXiv:2305.16317 (2023).\n\n[4] Durkan, Conor, et al. \"Neural spline flows.\" Advances in neural information processing systems 32 (2019).\n\n[5] Dax, Maximilian, et al. \"Real-time gravitational wave science with neural posterior estimation.\" Physical review letters 127.24 (2021): 241103.\n"}]

Author rebuttal:
We thank the reviewers for the positive evaluation of our work and for providing detailed comments and suggestions for improvement. In particular, all reviewers recommend acceptance to varying degrees, with the paper “very well written” (1iUK), “easy to understand” (7UFh), “offering meaningful contributions to the machine learning communities” (SM25) and with “nice theoretical results” (VmG3). There is general consensus that our presentation is clear and the work relevant to the field of Simulation-Based Inference (SBI).
The most common concern raised by reviewers relates to the perceived difficulty of our experiments and whether Flow Matching Posterior Estimation (FMPE) represents a substantial improvement over the state-of-the-art in SBI. In fact, our two-pronged approach, which involves a set of benchmark tests and a real-world problem, is designed to probe complementary aspects of the method, covering breadth and depth of applications:
1. **SBI benchmark [1]:** This establishes the applicability of FMPE across a wide range of SBI challenges. Indeed, the benchmark comprises a standard set of tests for which various methods (NPE, NRE, NLE, etc.) have been optimized. The results of Sec. 4 therefore allow “for easy comparisons against previous works in the literature” (1iUK). Some of these tests do involve toy models, but they have been chosen for reasons described in Sec. 2.3 of [1], e.g., to test how performance depends on dimensionality, complexity of posteriors, and the presence of distractors. We emphasize that the benchmark does contain some challenging tasks as well (Sec. 3, #2 in [1]: “This highlights that our problems &mdash; though conceptually simple &mdash; are challenging, and there is room for development of more powerful algorithms.”). Results on these tests show that FMPE is competitive with NPE, and also serve to validate our theoretical claims of probability-mass coverage.
2. **Gravitational wave inference:** This task is designed to push the performance limits of FMPE in a real-world problem. Indeed, this is arguably one of the most challenging scientific inference tasks for which rigorous quantitative comparisons are available, allowing for precise comparisons. It involves high-dimensional data (15744-dim, see l.315), complicated physics models, low signal-to-noise ratios, multi-detector measurements and potentially non-stationary noise in the real data. Furthermore, strict accuracy requirements must be met for results to be scientifically useful (l.301-302, see also [2]). In l.276, we cite eight previous studies applying SBI to GW inference (incl. NPE, NRE, variational inference), and (unlike FMPE) no other generic method (i.e., not leveraging additional symmetry information) meets these requirements.

To summarize, we believe that our experiments do indeed represent a thorough testing ground for FMPE, and that our method displays excellent performance. To update our paper, we will add a discussion that clarifies this reasoning behind our tests. As requested by reviewers, we will also clarify what we mean by "scalability" (it refers to data dimension) and provide additional details on limitations (e.g., slower inference than NPE). All other review comments are addressed in the individual responses.

[1] Lueckmann, Jan-Matthis, et al. "Benchmarking simulation-based inference." International conference on artificial intelligence and statistics. PMLR, 2021.


[2] Romero-Shaw, Isobel M., et al. "Bayesian inference for compact binary coalescences with bilby: validation and application to the first LIGO–Virgo gravitational-wave transient catalogue." Monthly Notices of the Royal Astronomical Society 499.3 (2020): 3295-3319.


Meta-review: The reviewers unanimously recommend acceptance (7-5-6-7). Minor issues have been raised and discussed during the author-reviewer discussion period. The authors are encouraged to take this feedback into account when preparing the final version of the paper.
Decision: Accept (poster)