SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jonas Wildberger*
Max Planck Institute for Intelligent Systems
wildberger.jonas@tuebingen.mpg.de
Simon Buchholz*
sbuchholz@tue.mpg.de
Jakob H. Macke
Max Planck Institute for Intelligent Systems &
Machine Learning in Science,
Maximilian Dax*
maximilian.dax@tuebingen.mpg.de
Stephen R. Green
Nottingham, United Kingdom
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.
The ability to readily represent Bayesian posteriors of arbitrary complexity using neural networks would herald a revolution in scientific data analysis. Such networks could be trained using simulated data and used for amortized inference across observations—bringing tractable inference and speed to a myriad of scientific models. Thanks to innovative architectures such as normalizing flows [1, 2], approaches to neural simulation-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.
Figure 1: Comparison of network architectures (left) and flow trajectories (right). Discrete flows (NPE, top) require a 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  $x$ , 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.
The Bayesian approach to data analysis is to compare observations to models via the posterior distribution  $p(\theta | x)$ . This gives our degree of belief that model parameters  $\theta$  gave rise to an observation  $x$ , and is proportional to the model likelihood  $p(x | \theta)$  times the prior  $p(\theta)$ . One is typically interested in representing the posterior in terms of a collection of samples, however obtaining these through standard likelihood-based algorithms can be challenging for intractable or expensive likelihoods. In such cases, SBI offers an alternative based instead on data simulations  $x \sim p(x | \theta)$ . Combined with deep generative modeling, SBI becomes a powerful paradigm for scientific inference [3]. Neural posterior estimation (NPE) [4-6], for instance, trains a conditional density estimator  $q(\theta | x)$  to approximate the posterior, allowing for rapid sampling and density estimation for any  $x$  consistent with the training distribution.
The NPE density estimator  $q(\theta | x)$  is commonly taken to be a (discrete) normalizing flow [1, 2], an approach that has brought state-of-the-art performance in challenging problems such as gravitational-wave inference [7]. Naturally, performance hinges on the expressiveness of  $q(\theta | x)$ . Normalizing flows transform noise to samples through a 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  $q(\theta | x)$ . Although many such discrete flows are universal density approximators [2], in practice, they can be challenging to scale to very large networks, which are needed for big-data experiments.
Recent studies [8, 9] propose neural posterior score estimation (NPSE), a rather different approach that models the posterior distribution with score-matching (or diffusion) networks. These techniques were originally developed for generative modeling [10-12], achieving state-of-the-art results in many domains, including image generation [13, 14]. Like discrete normalizing flows, diffusion models transform noise into samples, but with trajectories parametrized by a continuous "time" parameter  $t$ . The trajectories solve a stochastic differential equation [15] (SDE) defined in terms of a vector field  $v_{t}$ , which is trained to match the score of the intermediate distributions  $p_t$ . NPSE has several advantages compared to NPE, including the ability to combine multiple observations at inference time [9] 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 flow-matching posterior estimation (FMPE). Flow matching is also based on a vector field  $v_{t}$  and thereby also admits flexible network architectures (Fig. 1). For flow
matching, however,  $v_{t}$  directly defines the velocity field of sample trajectories, which solve ordinary differential equations (ODEs) and are deterministic. As a 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. 1.
NPE	NPSE	FMPE (Ours)
Tractable posterior density	Yes	No	Yes
Unconstrained network architecture	No	Yes	Yes
Network passes for sampling	Single	Many	Many
Table 1: Comparison of posterior-estimation methods.
- We adapt flow-matching to 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 a complex distribution (e.g., images). In contrast, for SBI,  $p(\theta | x)$  is typically less complex for fixed  $x$ , but  $x$  itself can be complex and high-dimensional. We therefore consider pyramid-like architectures from  $x$  to  $v_t$ , with gated linear units to incorporate  $(\theta, t)$  dependence, rather than the typical U-Net used for images (Fig. 1). 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 are masscovering, i.e., that their support includes all  $\theta$  consistent with observed  $x$ , 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 a 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-covering posteriors in 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 in terms of training time, posterior accuracy, and scaling to larger networks.
Normalizing flows. A normalizing flow [1, 2] defines a probability distribution  $q(\theta | x)$  over parameters  $\theta \in \mathbb{R}^n$  in terms of an invertible mapping  $\psi_x: \mathbb{R}^n \to \mathbb{R}^n$  from a simple base distribution  $q_0(\theta)$ ,
$$
q (\theta | x) = \left(\psi_ {x}\right) _ {*} q _ {0} (\theta) = q _ {0} \left(\psi_ {x} ^ {- 1} (\theta)\right) \det  \left| \frac {\partial \psi_ {x} ^ {- 1} (\theta)}{\partial \theta} \right|, \tag {1}
$$
where  $(\cdot)_*$  denotes the pushforward operator, and for generality we have conditioned on additional context  $x \in \mathbb{R}^m$ . Unless otherwise specified, a normalizing flow refers to a discrete flow, where  $\psi_x$  is given by a composition of simpler mappings with triangular Jacobians, interspersed with shuffling of the  $\theta$ . This construction results in expressive  $q(\theta|x)$  and also efficient density evaluation [2].
Continuous normalizing flows. A continuous flow [18] also maps from base to target distribution, but is parametrized by a continuous "time"  $t \in [0,1]$ , where  $q_0(\theta |x) = q_0(\theta)$  and  $q_{1}(\theta |x) = q(\theta |x)$ . For each  $t$ , the flow is defined by a vector field  $v_{t,x}$  on the sample space. This corresponds to the velocity of the sample trajectories,
$$
\frac {d}{d t} \psi_ {t, x} (\theta) = v _ {t, x} \left(\psi_ {t, x} (\theta)\right), \quad \psi_ {0, x} (\theta) = \theta . \tag {2}
$$
$^{2}$ Code available here.
<sup>3</sup>In the SBI literature, this is also commonly referred to as "parameter space".
We obtain the trajectories  $\theta_t \equiv \psi_{t,x}(\theta)$  by integrating this ODE. The final density is given by
$$
q (\theta | x) = \left(\psi_ {1, x}\right) _ {*} q _ {0} (\theta) = q _ {0} (\theta) \exp \left(- \int_ {0} ^ {1} \operatorname {d i v} v _ {t, x} \left(\theta_ {t}\right) \mathrm {d} t\right), \tag {3}
$$
which is obtained by solving the transport equation  $\partial_t q_t + \mathrm{div}(q_t v_{t,x}) = 0$ .
The advantage of the continuous flow is that  $v_{t,x}(\theta)$  can be simply specified by a neural network taking  $\mathbb{R}^{n + m + 1} \to \mathbb{R}^n$ , in which case (2) is referred to as a neural ODE [18]. Since the density is tractable via (3), it is in principle possible to train the flow by maximizing the (log-)likelihood. However, this is often not feasible in 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  $v_{t,x}$  on a vector field  $u_{t,x}$  that generates a target probability path  $p_{t,x}$ . It has the advantage that training does not require integration of ODEs, however it is not immediately clear how to choose  $(u_{t,x}, p_{t,x})$ . The key insight of [16] is that, if the path is chosen on a sample-conditional basis, then the training objective becomes extremely simple. Indeed, given a sample-conditional probability path  $p_t(\theta | \theta_1)$  and a corresponding vector field  $u_t(\theta | \theta_1)$ , we specify the sample-conditional flow matching loss as
$$
\mathcal {L} _ {\mathrm {S C F M}} = \mathbb {E} _ {t \sim \mathcal {U} [ 0, 1 ], x \sim p (x), \theta_ {1} \sim p (\theta | x), \theta_ {t} \sim p _ {t} (\theta_ {t} | \theta_ {1})} \| v _ {t, x} (\theta_ {t}) - u _ {t} (\theta_ {t} | \theta_ {1}) \| ^ {2}. \tag {4}
$$
Remarkably, minimization of this loss is equivalent to regressing  $v_{t,x}(\theta)$  on the marginal vector field  $u_{t,x}(\theta)$  that generates  $p_t(\theta |x)$  [16]. Note that in this expression, the  $x$ -dependence of  $v_{t,x}(\theta)$  is picked up via the expectation value, with the sample-conditional vector field independent of  $x$ .
There exists considerable freedom in choosing a sample-conditional path. Ref. [16] introduces the family of Gaussian paths
$$
p _ {t} (\theta | \theta_ {1}) = \mathcal {N} (\theta | \mu_ {t} (\theta_ {1}), \sigma_ {t} (\theta_ {1}) ^ {2} I _ {n}), \tag {5}
$$
where the time-dependent means  $\mu_t(\theta_1)$  and standard deviations  $\sigma_t(\theta_1)$  can be freely specified (subject to boundary conditions<sup>5</sup>). For our experiments, we focus on the optimal transport paths defined by  $\mu_t(\theta_1) = t\theta_1$  and  $\sigma_t(\theta_1) = 1 - (1 - \sigma_{\mathrm{min}})t$  (also introduced in [16]). The sample-conditional vector field then has the simple form
$$
u _ {t} \left(\theta \mid \theta_ {1}\right) = \frac {\theta_ {1} - \left(1 - \sigma_ {\min }\right) \theta}{1 - \left(1 - \sigma_ {\min }\right) t}. \tag {6}
$$
Neural posterior estimation (NPE). NPE is an SBI method that directly fits a density estimator  $q(\theta |x)$  (usually a normalizing flow) to the posterior  $p(\theta |x)$  [4-6]. NPE trains with the maximum likelihood objective  $\mathcal{L}_{\mathrm{NPE}} = -\mathbb{E}_{p(\theta)p(x|\theta)}\log q(\theta |x)$ , using Bayes' theorem to simplify the expectation value with  $\mathbb{E}_{p(x)p(\theta|x)}\to \mathbb{E}_{p(\theta)p(x|\theta)}$ . During training,  $\mathcal{L}_{\mathrm{NPE}}$  is estimated based on an empirical distribution consisting of samples  $(\theta ,x)\sim p(\theta)p(x|\theta)$ . Once trained, NPE can perform inference for every new observation using  $q(\theta |x)$ , thereby amortizing the computational cost of simulation and training across all observations. NPE further provides exact density evaluations of  $q(\theta |x)$ . Both of these properties are crucial for the physics application in section 5, so we aim to retain these properties with FMPE.
Flow matching [16] has been developed as a technique for generative modeling, and similar techniques are discussed in [19-21] and extended in [22, 23]. Flow matching encompasses the deterministic ODE version of diffusion models [10-12] as a special instance. Although to 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, use stochastic formulations via SDEs [15] or Langevin steps and are therefore not directly applicable when evaluations of the posterior density are desired (see Tab. 1). It should be noted that score modeling can also be used
<sup>4</sup>We refer to conditioning on  $\theta_{1}$  as sample-conditioning to distinguish from conditioning on  $x$ .
The sample-conditional probability path should be chosen to be concentrated around  $\theta_{1}$  at  $t = 1$  (within a small region of size  $\sigma_{\mathrm{min}}$ ) and to be the base distribution at  $t = 0$ .
to parameterize continuous normalizing flows via an ODE. Extension of [8, 9] to the deterministic formulation could thereby be seen as a special case of flow matching. Many of our analyses and the practical guidance provided in Section 3 therefore also apply to score matching.
We here focus on comparisons of FMPE against NPE [4-6], as it best matches the requirements of the application in section 5. Other SBI methods include approximate Bayesian computation [24-28], neural likelihood estimation [29-32] and neural ratio estimation [33-39]. Many of these approaches have sequential versions, where the estimator networks are specifically tuned to a specific observation  $x_{\mathrm{o}}$ . 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 to SBI we use Bayes' theorem to make the usual replacement  $\mathbb{E}_{p(x)p(\theta |x)}\to$ $\mathbb{E}_{p(\theta)p(x|\theta)}$  in the loss function (4), eliminating the intractable expectation values. This gives the FMPE loss
$$
\mathcal {L} _ {\mathrm {F M P E}} = \mathbb {E} _ {t \sim p (t), \theta_ {1} \sim p (\theta), x \sim p (x | \theta_ {1}), \theta_ {t} \sim p _ {t} (\theta_ {t} | \theta_ {1})} \| v _ {t, x} (\theta_ {t}) - u _ {t} (\theta_ {t} | \theta_ {1}) \| ^ {2}, \tag {7}
$$
which we minimize using empirical risk minimization over samples  $(\theta, x) \sim p(\theta)p(x|\theta)$ . In other words, training data is generated by sampling  $\theta$  from the prior, and then simulating data  $x$  corresponding to  $\theta$ . This is in close analogy to NPE training, but replaces the log likelihood maximization with the sample-conditional flow matching objective. Note that in this expression we also sample  $t \sim p(t)$ ,  $t \in [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(\theta | x)$  can achieve excellent results in approximating the true posterior  $p(\theta | x)$ . However, it 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  $q(\theta | x)$  is mass-covering, i.e., whether it contains the full support of  $p(\theta | x)$ . This minimizes the risk to falsely rule out possible explanations of the data. It also allows us to use importance sampling if the likelihood  $p(x | \theta)$  of the forward model can be evaluated, which can be used for precise estimation of the posterior [40, 41].
Consider first the mass-covering property for NPE. NPE directly minimizes the forward KL divergence  $\mathrm{D}_{\mathrm{KL}}(p(\theta |x)||q(\theta |x))$  and thereby provides probability-mass covering results. Therefore, even if NPE is not accurately trained, the estimate  $q(\theta |x)$  should cover the entire support of the posterior  $p(\theta |x)$  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 Gaussian  $q$  fitted to a bimodal target distribution  $p$  captures both modes when using the forward KL divergence  $\mathrm{D}_{\mathrm{KL}}(p||q)$ , but only a single mode when using the backwards direction  $\mathrm{D}_{\mathrm{KL}}(q||p)$  (Fig. 2).
For FMPE, we can fit a Gaussian flow-matching model  $q(\theta) = \mathcal{N}(\hat{\mu},\hat{\sigma}^2)$  to the same bimodal target, in this case, parametrizing the vector field as
Figure 2: A Gaussian (blue) fitted to a bimodal distribution (gray) with various objectives.
$$
v _ {t} (\theta) = \frac {\left(\sigma_ {t} ^ {2} + (t \hat {\sigma}) ^ {2} - \sigma_ {t}\right) \theta_ {t} + t \hat {\mu} \cdot \sigma_ {t}}{t \cdot \left(\sigma_ {t} ^ {2} + (t \hat {\sigma}) ^ {2}\right)} \tag {8}
$$
(see Appendix A), we also obtain a mass-covering distribution when fitting the learnable parameters  $(\hat{\mu},\hat{\sigma})$  via (4). 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
between vector fields  $u_{t}$  and  $v_{t}$  bounds the forward KL divergence. Indeed, the former agrees up to constant with the sample-conditional loss (4) (see Sec. 2).
We denote the flows of  $u_{t}, v_{t}$ , by  $\phi_t, \psi_t$ , respectively, and we set  $q_{t} = (\psi_{t})_{*}q_{0}$ ,  $p_{t} = (\phi_{t})_{*}q_{0}$ . The precise question then is whether we can bound  $\mathrm{D}_{\mathrm{KL}}(p_1||q_1)$  by  $\mathrm{MSE}_p(u,v)^\alpha$  for some positive power  $\alpha$ . It was already observed in [42] that this is not true in general, and we provide a simple example to that effect in Lemma 1 in Appendix B. Indeed, it was found in [42] that to bound the forward KL divergence we also need to control the Fisher divergence,  $\int p_t(\mathrm{d}\theta)(\nabla \ln p_t(\theta) - \nabla q_t(\theta))^2$ .
Here we show instead that a bound can be obtained under sufficiently strong regularity assumptions on  $p_0, u_t$ , and  $v_t$ . The following statement is slightly informal, and we refer to the supplement for the complete version.
Theorem 1. Let  $p_0 = q_0$  and assume  $u_t$  and  $v_t$  are two vector fields whose flows satisfy  $p_1 = (\phi_1)_*p_0$  and  $q_1 = (\psi_1)_*q_0$ . Assume that  $p_0$  is square integrable and satisfies  $|\nabla \ln p_0(\theta)| \leq c(1 + |\theta|)$  and  $u_t$  and  $v_t$  have bounded second derivatives. Then there is a constant  $C > 0$  such that (for  $\mathrm{MSE}_p(u,v) < 1$ )
$$
\mathrm {D} _ {\mathrm {K L}} \left(p _ {1} \mid \mid q _ {1}\right) \leq C \operatorname {M S E} _ {p} (u, v) ^ {\frac {1}{2}}. \tag {9}
$$
The proof of this result can be found in appendix B. While the regularity assumptions are not guaranteed to hold in practice when  $v_{t}$  is parametrized by a neural net, the theorem nevertheless gives some indication that the flow-matching objective encourages mass coverage. In Section 4 and 5, this is complemented with extensive empirical evidence that flow matching indeed provides mass-covering estimates.
We remark that it 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 Fokker-Planck equation in [42]. 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  $\theta$  space (e.g., images with millions of pixels). One typically uses U-Net [44] like architectures, as they provide a natural mapping from  $\theta$  to a vector field  $v(\theta)$  of the same dimension. The dependence on  $t$  and an (optional) conditioning vector  $x$  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  $\theta$  are typically low dimensional. In this context, it is therefore useful to build the architecture starting as a mapping from  $x$  to  $v(x)$  and then add conditioning on  $\theta$  and  $t$ . 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(\theta)$ . In our experiments, we found that the  $(t, \theta)$ -conditioning is best achieved using gated linear units [45] to the hidden layers of the network (see also Fig. 1); these are also commonly used for conditioning discrete flows on  $x$ .
3.3 Re-scaling the time prior
The time prior  $\mathcal{U}[0,1]$  in (4) distributes the training capacity uniformly across  $t$ . We observed that this is not always optimal in practice, as the complexity of the vector field may depend on  $t$ . For FMPE we therefore sample  $t$  in (7) from a power-law distribution  $p_{\alpha}(t)\propto t^{1 / (1 + \alpha)}$ ,  $t\in [0,1]$ , introducing an additional hyperparameter  $\alpha$ . This includes the uniform distribution for  $\alpha = 0$ , but for  $\alpha >0$ , assigns greater importance to the vector field for larger values of  $t$ . 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 in [46], ranging from simple Gaussian toy models to more challenging SBI problems from epidemiology and ecology, with varying dimensions for parameters  $(\dim(\theta) \in [2, 10])$  and observations  $(\dim(x) \in [2, 100])$ . For each task, we train three separate FMPE models with simulation budgets  $N \in \{10^3, 10^4, 10^5\}$ . We