SlowGuess/NeurIPS_2023_OCR huggingface dataset:
use a 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-Raw, SLCP-D), we condition on  $(t, \theta)$  via gated linear units as described in Section 3.2 (Fig. 8 in Appendix C shows the corresponding performance gain). For the remaining tasks with  $\dim(x) \leq 10$  we concatenate  $(t, \theta, x)$  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(\theta |x)$  to the reference posteriors  $p(\theta |x)$  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  $\theta \sim q(\theta |x)$  from reference samples  $\theta \sim p(\theta |x)$ . 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 an 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 a large simulation budget ( $N = 10^5$ ). 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 3.1. We display the density  $\log q(\theta |x)$  of the reference samples  $\theta \sim p(\theta |x)$  under our FMPE model  $q$  as a histogram (Fig. 3). All samples  $\theta \sim p(\theta |x)$  fall into the support from  $q(\theta |x)$ . This becomes apparent when comparing to the density  $\log q(\theta |x)$  for samples  $\theta \sim q(\theta |x)$  from  $q$  itself. This FMPE result is therefore mass covering. Note that this does not necessarily imply conservative posteriors (which is also not generally true for the forward KL divergence [17, 50, 51]), and some parts of  $p(\theta |x)$  may still be undersampled. Probability mass coverage, however, implies that no part is entirely missed (compare Fig. 2), even for multimodal distributions such as Two Moons. Fig. 9 in the Appendix confirms the mass coverage for the other benchmark tasks.
Figure 3: Histogram of FMPE densities  $\log q(\theta |x)$  for reference samples  $\theta \sim p(\theta |x)$  (Two Moons task,  $N = 10^{3}$ ). The estimate  $q(\theta |x)$  clearly covers  $p(\theta |x)$  entirely.
5 Gravitational-wave inference
5.1 Background
Gravitational waves (GWs) are ripples of spacetime predicted by Einstein and produced by cataclysmic cosmic events such as the mergers of binary black holes (BBHs). GWs propagate across the universe to Earth, where the LIGO-Virgo-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. 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
$^6$ BBH parameters  $\theta \in \mathbb{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 \in [20,512]\mathrm{Hz}$ ,  $\Delta f = 0.125\mathrm{Hz}$ , we have  $x \in \mathbb{R}^{15744}$ .
<sup>7</sup>Noise is assumed to be stationary and Gaussian, so for frequency-domain data, the GW likelihood  $p(x|\theta) = \mathcal{N}(h(\theta)|S_{\mathrm{n}})(x)$ . Here  $h(\theta)$  is a theoretical signal model based on Einstein's theory of general relativity, and  $S_{\mathrm{n}}$  is the power spectral density of the detector noise.
Figure 4: Comparison of FMPE with NPE, a standard SBI method, across 10 benchmark tasks [46].
typically performed using tools [57-60] based on Markov chain Monte Carlo [61, 62] or 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 to  $\sim 10^8$  likelihood evaluations. The ever-increasing rate of detections means that these analysis times risk becoming a bottleneck. SBI offers a promising solution for this challenge that has thus been actively studied in the literature [64-68, 7, 69, 70, 41]. A fully amortized NPE-based method called DINGO recently achieved accuracies comparable to stochastic samplers with inference times of less than a 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 a tractable density, which is problematic when verifying and correcting inference results using importance sampling [41].
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 a few minor changes (see Appendix D). This uses an 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 a learnable linear layer initialized with principal components of GW simulations followed by a series of dense residual blocks [47]. This architecture is a 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,\theta)$  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  $\theta$  so that the network parameterizes the conditional vector field  $(t,x,\theta)\rightarrow v_{t,x}(\theta)$ ; (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 \cdot 10^{6}$  simulations for 400 epochs using a batch size of 4096 on an A100 GPU. The FMPE network ( $1.9 \cdot 10^{8}$  learnable parameters, training takes  $\approx 2$  days) is larger than the NPE network ( $1.3 \cdot 10^{8}$  learnable parameters, training takes  $\approx 3$  days), but trains substantially faster. We evaluate both networks on GW150914 [73], the first detected GW. We generate a reference posterior using the method described in [41]. Fig. 5 compares the inferred posterior distributions qualitatively and quantitatively in terms of the Jensen-Shannon divergence (JSD) to the reference. $^{9}$
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 indistinguishability 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://github.com/dingo-gw/dingo.
<sup>9</sup>We omit the three parameters  $\phi_{\mathrm{c}}, \phi_{JL}, \theta_{JN}$  in the evaluation as we use phase marginalization in importance sampling and the reference therefore uses a different basis for these parameters [41]. For GNPE we report the results from [7], 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 a rough comparison. The JSD for the  $t_{\mathrm{c}}$  parameter is not reported in [7] due to a  $t_{\mathrm{c}}$  marginalized reference.
Figure 5: Results for GW150914 [73]. Left: Corner plot showing 1D marginals on the diagonal and 2D  $50\%$  credible regions. We display four GW parameters (distance  $d_{\mathrm{L}}$ , time of arrival  $t_{\mathrm{c}}$ , and sky coordinates  $\alpha, \delta$ ); these represent the least accurate NPE parameters. Right: Deviation between inferred posteriors and the reference, quantified by the Jensen-Shannon divergence (JSD). 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, we find that the Bayesian evidences inferred with NPE  $(\log p(x) = -7667.958 \pm 0.006)$  and FMPE  $(\log p(x) = -7667.969 \pm 0.005)$  are consistent within their statistical uncertainties. A correct evidence is only obtained in importance sampling when the inferred posterior  $q(\theta | x)$  covers the entire posterior  $p(\theta | x)$  [41], so this is another indication that FMPE indeed induces mass-covering posteriors.
Our results for GW150914 show that FMPE substantially outperforms NPE on this challenging problem. We believe that this is related to the network structure as 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 of the high-dimensional  $x \in \mathbb{R}^{15744}$ . Hence, it scales better to larger networks and achieves higher accuracy. Remarkably, the performance is comparable to GNPE, which involves a 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.
In future work we plan to carry out a more complete analysis of GW inference using FMPE. Indeed, GW150914 is a loud event with good data quality, where NPE already performs quite well. DINGO with GNPE has been validated in a variety of settings [7, 69, 41, 74] including events with a 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. A straightforward application of the GNPE mechanism to FMPE—GFMPE—is also possible, but less practical due to the higher inference costs of FMPE. Nevertheless, our results demonstrate that FMPE is a promising direction for future research in this field.
We introduced flow matching posterior estimation, a new simulation-based 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 a 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 an 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/backward passes.
We evaluated FMPE on a set of 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 a method that explicitly leverages symmetries to simplify training. Additionally, flow matching latent spaces are more naturally structured than those of discrete flows, particularly when using paths such as optimal transport. Looking forward, it would be interesting to exploit such structure in designing learning algorithms. This performance and flexibility underscores the capability of continuous normalizing flows to efficiently solve inverse problems.
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 is 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-Society (MPS), and the State of Niedersachsen/Germany 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 is 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 contributions 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 of 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 01IS18039A).
[1] Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International Conference on Machine Learning, pages 1530-1538, 2015. arXiv:1505.05770.
[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):1-64, 2021. URL: http://jmlr.org/papers/v22/19-1028.html.
[3] Kyle Cranmer, Johann Brehmer, and Gilles Louppe. The frontier of simulation-based inference. Proc. Nat. Acad. Sci., 117(48):30055-30062, 2020. arXiv:1911.01429, doi:10.1073/pnas.1912789117.
[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.06281.
[72] Conor Durkan, Artur Bekasov, Iain 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 et al. Observation of Gravitational Waves from a Binary Black Hole Merger. Phys. Rev. Lett., 116(6):061102, 2016. arXiv:1602.03837, doi:10.1103/PhysRevLett.116.061102.
[74] Jonas Wildberger, Maximilian Dax, Stephen R. Green, Jonathan Gair, Michael Purrer, 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/PhysRevD.107.084046.
[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-2999. JMLR.org, 2016. URL: http://proceedings.mlr.press/v48/cohenc16.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. Lett., 113:151101, Oct 2014. URL: https://link.aps.org/doi/10.1103/PhysRevLett.113.151101, doi:10.1103/PhysRevLett.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. II. A phenomenological model for the advanced detector era. Phys. Rev., D93(4):044007, 2016. arXiv:1508.07253, doi:10.1103/PhysRevD.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-T1500602-v4, 2016. URL: https://dcc.ligo.org/LIGO-T1500602/public.
[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:10.1103/PhysRevD.90.024018.
[80] J.R. Dormand and P.J. Prince. A family of embedded runge-kutta formulae. Journal of Computational and Applied Mathematics, 6(1):19-26, 1980. URL: https:// www.sciencedirect.com/science/article/pii/0771050X80900133, doi:https:// doi.org/10.1016/0771-050X(80)90013-3.
A Gaussian flow
We here derive the form of a vector field  $v_{t}(\theta)$  that restricts the resulting continuous flow to a one-dimensional Gaussian with mean  $\hat{\mu}$  variance  $\hat{\sigma}^2$ . With the optimal transport path  $\mu_t(\theta) = t\theta_1$ ,  $\sigma_t(\theta) = 1 - (1 - \sigma_{\mathrm{min}})t \equiv \sigma_t$  from [16], the sample-conditional probability path (5) reads
$$
p _ {t} (\theta | \theta_ {1}) = \mathcal {N} [ t \theta_ {1}, \sigma_ {t} ^ {2} ] (\theta). \tag {10}
$$
We set our target distribution
$$
q _ {1} \left(\theta_ {1}\right) = \mathcal {N} \left[ \hat {\mu}, \hat {\sigma} ^ {2} \right] \left(\theta_ {1}\right). \tag {11}
$$
To derive the marginal probability path and the marginal vector field we need two identities for the convolution  $*$  of Gaussian densities. Recall that the convolution of two functions is defined by  $f*g(x) = \int f(x - y)g(y)\mathrm{d}y$ . We define the function
$$
g _ {\mu , \sigma^ {2}} (\theta) = \theta \cdot \mathcal {N} [ \mu , \sigma^ {2} ] (\theta). \tag {12}
$$
Then the following holds
$$
\mathcal {N} \left[ \mu_ {1}, \sigma_ {1} ^ {2} \right] * \mathcal {N} \left[ \mu_ {2}, \sigma_ {2} ^ {2} \right] = \mathcal {N} \left[ \mu_ {1} + \mu_ {2}, \sigma_ {1} ^ {2} + \sigma_ {2} ^ {2} \right] \tag {13}
$$
$$
\left. g _ {0, \sigma_ {1} ^ {2}} * \mathcal {N} \left[ \mu_ {2}, \sigma_ {2} ^ {2} \right] = \frac {\sigma_ {1} ^ {2}}{\sigma_ {1} ^ {2} + \sigma_ {2} ^ {2}} \left(g _ {\mu_ {2}, \sigma_ {1} ^ {2} + \sigma_ {2} ^ {2}} - \mu_ {2} \mathcal {N} \left[ \mu_ {2}, \sigma_ {1} ^ {2} + \sigma_ {2} ^ {2} \right]\right) \right. \tag {14}
$$
Marginal probability paths
Marginalizing over  $\theta_{1}$  in (10) with (11), we find
$$
\begin{array}{l} p _ {t} (\theta) = \int p _ {t} (\theta | \theta_ {1}) q (\theta_ {1}) d \theta_ {1} \\ = \int \mathcal {N} \left[ t \theta_ {1}, \sigma_ {t} ^ {2} \right] (\theta) \mathcal {N} \left[ \hat {\mu}, \hat {\sigma} ^ {2} \right] (\theta_ {1}) d \theta_ {1} \\ = \int \mathcal {N} \left[ 0, \sigma_ {t} ^ {2} \right] (\theta - t \theta_ {1}) \mathcal {N} \left(t \hat {\mu}, (t \hat {\sigma}) ^ {2}\right) (t \theta_ {1}) \cdot t d \theta_ {1} \tag {15} \\ = \int \mathcal {N} \left[ 0, \sigma_ {t} ^ {2} \right] (\theta - \theta_ {1} ^ {t}) \mathcal {N} \left[ t \hat {\mu}, (t \hat {\sigma}) ^ {2} \right] (\theta_ {1} ^ {t}) d \theta_ {1} ^ {t} \\ = \mathcal {N} \left[ t \hat {\mu}, \sigma_ {t} ^ {2} + (t \hat {\sigma}) ^ {2} \right] (\theta) \\ \end{array}
$$
where we defined  $\theta_1^t = t\theta_1$  and used (13).
Marginal vector field
We now calculate the marginalized vector field  $u_{t}(\theta)$  based on equation (8) in [16]. Using the sample-conditional vector field (6) and the distributions (10) and (11) we find
$$
\begin{array}{l} u _ {t} (\theta) = \int u _ {t} (\theta | \theta_ {1}) \frac {p _ {t} (\theta | \theta_ {1}) q (\theta_ {1})}{p _ {t} (\theta)} d \theta_ {1} \\ = \frac {1}{p _ {t} (\theta)} \int \frac {\left(\theta_ {1} - \left(1 - \sigma_ {\min }\right) \theta\right)}{\sigma_ {t}} \cdot \mathcal {N} \left[ t \theta_ {1}, \sigma_ {t} ^ {2} \right] (\theta) \cdot \mathcal {N} \left[ \hat {\mu}, \hat {\sigma} ^ {2} \right] (\theta_ {1}) d \theta_ {1} \\ = \frac {1}{p _ {t} (\theta)} \int \frac {\left(\theta_ {1} - \left(1 - \sigma_ {\min }\right) \theta\right)}{\sigma_ {t}} \cdot \mathcal {N} \left[ 0, \sigma_ {t} ^ {2} \right] (\theta - t \theta_ {1}) \cdot \mathcal {N} \left[ t \hat {\mu}, (t \hat {\sigma}) ^ {2} \right] (t \theta_ {1}) \cdot t d \theta_ {1} \\ = \frac {1}{p _ {t} (\theta)} \int \frac {\left(\theta_ {1} ^ {\prime} - (1 - \sigma_ {\min }) t \cdot \theta\right)}{\sigma_ {t} \cdot t} \cdot \mathcal {N} \left[ 0, \sigma_ {t} ^ {2} \right] (\theta - \theta_ {1} ^ {\prime}) \cdot \mathcal {N} \left[ t \hat {\mu}, (t \hat {\sigma}) ^ {2} \right] (\theta_ {1} ^ {\prime}) \cdot d \theta_ {1} ^ {\prime} \\ = \frac {1}{p _ {t} (\theta)} \int \frac {\left(- \theta_ {1} ^ {\prime \prime} + (1 - (1 - \sigma_ {\min }) t) \cdot \theta\right)}{\sigma_ {t} \cdot t} \cdot \mathcal {N} \left[ 0, \sigma_ {t} ^ {2} \right] \left(\theta_ {1} ^ {\prime \prime}\right) \cdot \mathcal {N} \left[ t \hat {\mu}, (t \hat {\sigma}) ^ {2} \right] (\theta - \theta_ {1} ^ {\prime \prime}) \cdot d \theta_ {1} ^ {\prime \prime} \\ = \frac {1}{p _ {t} (\theta)} \int \frac {\left(- \theta_ {1} ^ {\prime \prime} + \sigma_ {t} \cdot \theta\right)}{\sigma_ {t} \cdot t} \cdot \mathcal {N} \left[ 0, \sigma_ {t} ^ {2} \right] \left(\theta_ {1} ^ {\prime \prime}\right) \cdot \mathcal {N} \left[ t \hat {\mu}, (t \hat {\sigma}) ^ {2} \right] (\theta - \theta_ {1} ^ {\prime \prime}) \cdot d \theta_ {1} ^ {\prime \prime} \tag {16} \\ \end{array}
$$
where we used the change of variables  $\theta_1' = t\theta_1$  and  $\theta_1'' = \theta - \theta_1'$ . Now we evaluate this expression using (12), then the identities (13) and (14) and the marginal probability (15)