Title: Flow Matching for Scalable Simulation-Based Inference
Venue: NeurIPS 2023
Accepted: yes
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-Based Inference

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.

1 Introduction

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.

Bernhard Schölkopf

<latexi sh

1_b

64="pm qZ+A

fOv/2H

P

>

B

3

c

VDLSgN

E

y

UY9

o

K

wG

8

T

k

WM

u

zC

J

j

R

5

X

d

7

<latexi sh

1_b

64="/Efv

n

3T

g

98

p

zVY

Z

>A

B

c

DLS

N

Oy

r

U

oP

K

wG

J

H

5I

2W

m

0

+kX

R

C

q

u

d

M

j

F

Q

7

<latexi sh

1_b

64="7

RI

qj

ZvJmX

/S

dHg

5

>A

B

c

VDL

N

EOy

r

fUY9

oP

K

M

w

W

p

F

Q

u

2

G

z3

8

+

C

k

0

n

T

<latexi sh

1_b

64="7

RI

qj

ZvJmX

/S

dHg

5

>A

B

c

VDL

N

EOy

r

fUY9

oP

K

M

w

W

p

F

Q

u

2

G

z3

8

+

C

k

0

n

T

<latexi sh

1_b

64="w/T

ud

NQ

PzIV

JY

Xg

F

>A

B9H

c

D

S

EO2

W

K

k

j

y

GnU5M0

r

m

v

f

3

p

L

R

7

o

Z

q

C

+

8

<latexi sh

1_b

64="

F

o

J

5cdp3RQY

M

>A

B+n

VD

S

PW

Eq

0

K

9

y

g

u

/

z

7

k

w

2v

O

8

m

U

j

C

I

X

f

H

<latexi sh

1_b

64="DLw

0B

Amc

/E

n

CyQKH

U

>

+

V

S

N

FJ3

rPW

Y

q

o

9

p

M

T

g

u

z

5d

7

kR

2v

O

8

G

j

Z

I

X

f

<latexi sh

1_b

64="Dpdu/

0ZH

PMJ

n

Q

o

v

G

>A

c

V

LS

N

F

3

r

W

Y

Eq

K

9

y

T

g

z

5

7

kRw

2

O

8

m

U

j

C

I

X

f

<latexi sh

1_b

64="

G5VY

Q2vm/K

9uOy

Nq

I

>A

CD3

c

LS

FJ

U

E

g0Wp

RX

P

T

f

M

B

j

8

o

d

7

w

Z

r

k

n

z

H

+

<latexi sh

1_b

64="dr

YELvgP

H8

O9F

>A

B

3

VDJS

N

2

W

5

R

M

I

TqU

C

zG

/

j

o

Z

7

f

un

0

w

Q

k

y+

<latexi sh

1_b

64="7

p

kgTBLDJPR

jG

9

Z2COWU

>A

+n

c

V

S

N

F

3

r

Y

Eq

5

y

0I

m

o

u

/

z

Q

HM

v

d

w

8

f

K

X

<latexi sh

1_b

64="8BuNO7

AM/

K

w+V

R

q

dSIc

>

C

X

DJ

g

E

2

W

v

pT

m

0k

5

F

9

j

P

H

Q

G

oU3

r

Ly

f

z

Y

n

Z

<latexi sh

1_b

64="8BuNO7

AM/

K

w+V

R

q

dSIc

>

C

X

DJ

g

E

2

W

v

pT

m

0k

5

F

9

j

P

H

Q

G

oU3

r

Ly

f

z

Y

n

Z

q

(

θ

|

x

o

)

ψ

i

(

θ

, x

)

ψ

n

(

θ

, x

)

ψ

1

(

θ

, x

)

θ

∼

N

0

,I

(

θ

)

q

(

θ

|

x

)

x

q

(

θ

|

x

o

)

x

(

θ

t

, t

)

v

t,x

(

θ

t

)

FMPE

NPE

The Bayesian approach to data analysis is to compare observations to models via the posterior distribution p ( θ | x ) . This gives our degree of belief that model parameters θ gave rise to an observation x , and is proportional to the model likelihood p ( x | θ ) times the prior p ( θ ) . 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 ∼ p ( x | θ ) . 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 ( θ | 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 ( θ | 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 gravitationalwave inference [7]. Naturally, performance hinges on the expressiveness of q ( θ | 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 ( θ | 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.

Our contributions are as follows:

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 ( θ | 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 ( θ, 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 θ consistent with observed x , which is highly desirable for scientific applications [17].

We perform a number of experiments to investigate the performance of FMPE. 2 Our twopronged 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.

2 Preliminaries

Normalizing flows. A normalizing flow [1, 2] defines a probability distribution q ( θ | x ) over parameters θ ∈ R n in terms of an invertible mapping ψ x : R n → R n from a simple base distribution q 0 ( θ ) ,

where ( · ) ∗ denotes the pushforward operator, and for generality we have conditioned on additional context x ∈ R m . Unless otherwise specified, a normalizing flow refers to a discrete flow, where ψ x is given by a composition of simpler mappings with triangular Jacobians, interspersed with shuffling of the θ . This construction results in expressive q ( θ | 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 ∈ [0 , 1] , where q 0 ( θ | x ) = q 0 ( θ ) and q 1 ( θ | x ) = q ( θ | x ) . For each t , the flow is defined by a vector field v t,x on the sample space. 3 This corresponds to the velocity of the sample trajectories,

We obtain the trajectories θ t ≡ ψ t,x ( θ ) by integrating this ODE. The final density is given by

which is obtained by solving the transport equation ∂ t q t +div( q t v t,x ) = 0 .

The advantage of the continuous flow is that v t,x ( θ ) can be simply specified by a neural network taking R n + m +1 → 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, 4 then the training objective becomes extremely simple. Indeed, given a sample-conditional probability path p t ( θ | θ 1 ) and a corresponding vector field u t ( θ | θ 1 ) , we specify the sample-conditional flow matching loss as

Remarkably, minimization of this loss is equivalent to regressing v t,x ( θ ) on the marginal vector field u t,x ( θ ) that generates p t ( θ | x ) [16]. Note that in this expression, the x -dependence of v t,x ( θ ) 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 where the time-dependent means µ t ( θ 1 ) and standard deviations σ t ( θ 1 ) can be freely specified (subject to boundary conditions 5 ). For our experiments, we focus on the optimal transport paths defined by µ t ( θ 1 ) = tθ 1 and σ t ( θ 1 ) = 1 -(1 -σ min ) t (also introduced in [16]). The sampleconditional vector field then has the simple form

Neural posterior estimation (NPE). NPE is an SBI method that directly fits a density estimator q ( θ | x ) (usually a normalizing flow) to the posterior p ( θ | x ) [4-6]. NPE trains with the maximum likelihood objective L NPE = -E p ( θ ) p ( x | θ ) log q ( θ | x ) , using Bayes' theorem to simplify the expectation value with E p ( x ) p ( θ | x ) → E p ( θ ) p ( x | θ ) . During training, L NPE is estimated based on an empirical distribution consisting of samples ( θ, x ) ∼ p ( θ ) p ( x | θ ) . Once trained, NPE can perform inference for every new observation using q ( θ | x ) , thereby amortizing the computational cost of simulation and training across all observations. NPE further provides exact density evaluations of q ( θ | x ) . Both of these properties are crucial for the physics application in section 5, so we aim to retain these properties with FMPE.

Related work

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

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 o . FMPEhas 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 E p ( x ) p ( θ | x ) → E p ( θ ) p ( x | θ ) in the loss function (4), eliminating the intractable expectation values. This gives the FMPE loss

which we minimize using empirical risk minimization over samples ( θ, x ) ∼ p ( θ ) p ( x | θ ) . In other words, training data is generated by sampling θ from the prior, and then simulating data x corresponding to θ . 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 ∼ p ( t ) , t ∈ [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 ( θ | x ) can achieve excellent results in approximating the true posterior p ( θ | 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 ( θ | x ) is mass-covering , i.e., whether it contains the full support of p ( θ | 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 | θ ) 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 D KL ( p ( θ | x ) || q ( θ | x )) , and thereby provides probabilitymass covering results. Therefore, even if NPE is not accurately trained, the estimate q ( θ | x ) should cover the entire support of the posterior p ( θ | 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 D KL ( p || q ) , but only a single mode when using the backwards direction D KL ( q || p ) (Fig. 2).

argmin

q

D

KL

(

q

||

p

)

argmin

q

D

KL

(

p

||

q

)

argmin

q

L

FM

For FMPE, we can fit a Gaussian flow-matching model q ( θ ) = N (ˆ µ, ˆ σ 2 ) to the same bimodal target, in this case, parametrizing the vector field as

(see Appendix A), we also obtain a mass-covering distribution when fitting the learnable parameters (ˆ µ, ˆ σ ) 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 ϕ t , ψ t , respectively, and we set q t = ( ψ t ) ∗ q 0 , p t = ( ϕ t ) ∗ q 0 . The precise question then is whether we can bound D KL ( p 1 || q 1 ) by MSE p ( u, v ) α for some positive power α . 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, ∫ p t (d θ )( ∇ ln p t ( θ ) -∇ q t ( θ )) 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 = ( ϕ 1 ) ∗ p 0 and q 1 = ( ψ 1 ) ∗ q 0 . Assume that p 0 is square integrable and satisfies |∇ ln p 0 ( θ ) | ≤ c (1 + | θ | ) and u t and v t have bounded second derivatives. Then there is a constant C > 0 such that (for MSE p ( u, v ) < 1) )

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 θ space (e.g., images with millions of pixels). One typically uses U-Net [44] like architectures, as they provide a natural mapping from θ to a vector field v ( θ ) 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 θ 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 θ 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( θ ) . In our experiments, we found that the ( t, θ ) -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 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 α ( t ) ∝ t 1 / (1+ α ) , t ∈ [0 , 1] , introducing an additional hyperparameter α . This includes the uniform distribution for α = 0 , but for α > 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( θ ) ∈ [2 , 10] ) and observations ( dim( x ) ∈ [2 , 100] ). For each task, we train three separate FMPE models with simulation budgets N ∈ { 10 3 , 10 4 , 10 5 } . We

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, θ ) 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 ) ≤ 10 we concatenate ( t, θ, 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 ( θ | x ) to the reference posteriors p ( θ | 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 θ ∼ q ( θ | x ) from reference samples θ ∼ p ( θ | 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 ( θ | x ) of the reference samples θ ∼ p ( θ | x ) under our FMPE model q as a histogram (Fig. 3). All samples θ ∼ p ( θ | x ) fall into the support from q ( θ | x ) . This becomes apparent when comparing to the density log q ( θ | x ) for samples θ ∼ q ( θ | 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 ( θ | 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.

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. 6 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, 7 and inference is

-

4

-

2

0

2

4

log

q

(

θ

x

)

θ

∼

p

(

θ

|

x

)

θ

∼

q

(

θ

|

x

)

|

0

.

6

0

.

8

1

.

0

GL

GL-U

GM

Two Moons

SLCP

10

3

10

4

10

5

0

.

6

0

.

8

1

.

0

B-GLM

10

3

10

4

10

5

B-GLM-Raw

10

3

10

4

10

5

SLCP-D

10

3

10

4

10

5

SIR

10

3

10

4

10

5

LV

Number of Simulations

C2ST

NPE

FMPE

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 ∼ 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). 8 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, θ ) 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 θ so that the network parameterizes the conditional vector field ( t, x, θ ) → v t,x ( θ ) ; (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 6 simulations for 400 epochs using a batch size of 4096 on an A100 GPU. The FMPE network ( 1 . 9 · 10 8 learnable parameters, training takes ≈ 2 days) is larger than the NPE network ( 1 . 3 · 10 8 learnable parameters, training takes ≈ 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

Reference

NPE

FMPE

-

8

0

8

t

c

[ms]

1

.

0

1

.

5

2

.

0

2

.

5

α

200

400

600

d

L

-

1

.

0

-

0

.

5

0

.

0

δ

-

8

0

8

t

c

[ms]

1

.

0

1

.

5

2

.

0

2

.

5

α

-

1

.

0

-

0

.

5

0

.

0

δ

NPE

FMPE

GNPE

m

1

m

2

a

1

a

2

t

1

t

2

φ

12

d

L

t

c

α

δ

ψ

1.2

1.3

0.8

2.5

0.6

1.1

3.2

0.8

0.2

1.6

1.0

0.3

0.8

0.1

0.5

0.4

0.3

0.5

0.3

0.2

0.1

4.4

0.1

0.8

9.1

0.6

-

10.1

0.6

0.7

8.6

0.5

1.4

0.6

0.1

0.2

0

5

10

15

20

JSD

[mnat]

to simplify data. Finally, we find that the Bayesian evidences inferred with NPE ( log p ( x ) = -7667 . 958 ± 0 . 006 ) and FMPE ( log p ( x ) = -7667 . 969 ± 0 . 005 ) are consistent within their statistical uncertainties. A correct evidence is only obtained in importance sampling when the inferred posterior q ( θ | x ) covers the entire posterior p ( θ | x ) [41], so this is 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 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 ∈ R 15744 . Hence, it scales better to larger networks and achieves higher accuracy. Remarkably, the performance iscomparable 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.

6 Conclusions

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 flexibilty underscores the capability of continuous normalizing flows to efficiently solve inverse problems.

References

Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International Conference on Machine Learning , pages 1530-1538, 2015. arXiv:1505.05770 .

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 .

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 .

George Papamakarios and Iain Murray. Fast ε -free inference of simulation models with Bayesian conditional density estimation. In Advances in neural information processing systems , 2016. arXiv:1605.06376 .

Jan-Matthis Lueckmann, Pedro J Gonçalves, Giacomo Bassetto, Kaan Öcal, Marcel Nonnenmacher, and Jakob H Macke. Flexible statistical inference for mechanistic models of neural dynamics. In Proceedings of the 31st International Conference on Neural Information Processing Systems , pages 1289-1299, 2017.

David Greenberg, Marcel Nonnenmacher, and Jakob Macke. Automatic posterior transformation for likelihood-free inference. In International Conference on Machine Learning , pages 24042414. PMLR, 2019.

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):241103, 2021. arXiv:2106.12594 , doi:10.1103/PhysRevLett. 127.241103 .

Louis Sharrock, Jack Simons, Song Liu, and Mark Beaumont. Sequential neural score estimation: Likelihood-free inference with conditional score based diffusion models. arXiv preprint arXiv:2210.04872 , 2022.

Tomas Geffner, George Papamakarios, and Andriy Mnih. Score modeling for simulation-based inference. arXiv preprint arXiv:2209.14249 , 2022.

Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning , pages 2256-2265. PMLR, 2015.

Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems , 32, 2019.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems , 33:6840-6851, 2020.

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://proceedings.neurips.cc/paper_files/paper/ 2021/file/49ad23d1ec9fa4bd8d77d02681df5cfa-Paper.pdf .

Jonathan Ho, Chitwan Saharia, William Chan, David J. 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: http://jmlr.org/papers/v23/21-0635.html .

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.13456 , 2020.

Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. CoRR , abs/2210.02747, 2022. arXiv:2210.02747 , doi: 10.48550/arXiv.2210.02747 .

Joeri Hermans, Arnaud Delaunoy, François Rozet, Antoine Wehenkel, and Gilles Louppe. Averting a crisis in simulation-based inference. arXiv preprint arXiv:2110.06581 , 2021.

Tian Qi Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa-Bianchi, and Roman Garnett, editors, Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurIPS 2018, December 3-8, 2018, Montréal, Canada , pages 6572-6583, 2018. URL: https://proceedings.neurips.cc/paper/2018/hash/ 69386f6bb1dfed68692a24c8686939b9-Abstract.html .

Michael S Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants. arXiv preprint arXiv:2209.15571 , 2022.

Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.03003 , 2022.

Kirill Neklyudov, Daniel Severo, and Alireza Makhzani. Action matching: A variational method for learning stochastic dynamics from samples. arXiv preprint arXiv:2210.06662 , 2022.

Aram Davtyan, Sepehr Sameni, and Paolo Favaro. Randomized conditional flow matching for video prediction. arXiv preprint arXiv:2211.14575 , 2022.

Alexander Tong, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, Kilian Fatras, Guy Wolf, and Yoshua Bengio. Conditional flow matching: Simulation-free dynamic optimal transport. arXiv preprint arXiv:2302.00482 , 2023.

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.

Mark A Beaumont, Wenyang Zhang, and David J Balding. Approximate bayesian computation in population genetics. Genetics , 162(4):2025-2035, 2002.

Mark A Beaumont, Jean-Marie Cornuet, Jean-Michel Marin, and Christian P Robert. Adaptive approximate bayesian computation. Biometrika , 96(4):983-990, 2009.

Michael G. B. Blum and Olivier François. Non-linear regression models for approximate bayesian computation. Stat. Comput. , 20(1):63-73, 2010. doi:10.1007/ s11222-009-9116-0 .

Dennis Prangle, Paul Fearnhead, Murray P. Cox, Patrick J. Biggs, and Nigel P. French. Semiautomatic selection of summary statistics for abc model choice, 2013. arXiv:1302.5624 .

Simon Wood. Statistical inference for noisy nonlinear ecological dynamic systems. Nature , 466:1102-4, 08 2010. doi:10.1038/nature09319 .

Christopher C Drovandi, Clara Grazian, Kerrie Mengersen, and Christian Robert. Approximating the likelihood in approximate bayesian computation, 2018. arXiv:1803.06645 .

George Papamakarios, David Sterratt, and Iain Murray. Sequential neural likelihood: Fast likelihood-free inference with autoregressive flows. In The 22nd International Conference on Artificial Intelligence and Statistics , pages 837-848. PMLR, 2019.

Jan-Matthis Lueckmann, Giacomo Bassetto, Theofanis Karaletsos, and Jakob H Macke. Likelihood-free inference with emulator networks. In Symposium on Advances in Approximate Bayesian Inference , pages 32-53. PMLR, 2019.

Rafael Izbicki, Ann Lee, and Chad Schafer. High-dimensional density ratio estimation with extensions to approximate likelihood computation. In Artificial intelligence and statistics , pages 420-429. PMLR, 2014.

Kim Pham, David Nott, and Sanjay Chaudhuri. A note on approximating abc-mcmc using flexible classifiers. Stat , 3, 03 2014. doi:10.1002/sta4.56 .

Kyle Cranmer, Juan Pavez, and Gilles Louppe. Approximating likelihood ratios with calibrated discriminative classifiers. arXiv preprint arXiv:1506.02169 , 2015.

Joeri Hermans, Volodimir Begy, and Gilles Louppe. Likelihood-free mcmc with approximate likelihood ratios. arXiv preprint arXiv:1903.04057 , 10, 2019.

Conor Durkan, Iain Murray, and George Papamakarios. On contrastive learning for likelihoodfree inference. In International conference on machine learning , pages 2771-2781. PMLR, 2020.

Owen Thomas, Ritabrata Dutta, Jukka Corander, Samuel Kaski, and Michael U. Gutmann. Likelihood-free inference by ratio estimation, 2020. arXiv:1611.10242 .

Benjamin K Miller, Christoph Weniger, and Patrick Forré. Contrastive neural ratio estimation. Advances in Neural Information Processing Systems , 35:3262-3278, 2022.

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.

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-Wave Inference. Phys. Rev. Lett. , 130(17):171403, 2023. arXiv:2210.05686 , doi:10.1103/PhysRevLett.130.171403 .

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.

Yang Song, Conor Durkan, Iain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. Advances in Neural Information Processing Systems , 34:14151428, 2021.

Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical Image Computing and Computer-Assisted Intervention-MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18 , pages 234-241. Springer, 2015.

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.

Jan-Matthis Lueckmann, Jan Boelts, David Greenberg, Pedro Goncalves, and Jakob Macke. Benchmarking simulation-based inference. In International Conference on Artificial Intelligence and Statistics , pages 343-351. PMLR, 2021.

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition, 2015. arXiv:1512.03385 .

Jerome H Friedman. On multivariate goodness-of-fit and two-sample testing. Statistical Problems in Particle Physics, Astrophysics, and Cosmology , 1:311, 2003.

David Lopez-Paz and Maxime Oquab. Revisiting classifier two-sample tests. arXiv preprint arXiv:1610.06545 , 2016.

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 .

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.

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 , 11 2021. arXiv:2111.03606 .

R. Abbott et al. Population Properties of Compact Objects from the Second LIGO-Virgo Gravitational-Wave Transient Catalog. Astrophys. J. Lett. , 913(1):L7, 2021. arXiv:2010. 14533 , doi:10.3847/2041-8213/abe949 .

B. P. Abbott et al. GW170817: Measurements of neutron star radii and equation of state. Phys. Rev. Lett. , 121(16):161101, 2018. arXiv:1805.11581 , doi:10.1103/PhysRevLett.121. 161101 .

R. Abbott et al. Tests of general relativity with binary black holes from the second LIGO-Virgo gravitational-wave transient catalog. Phys. Rev. D , 103(12):122002, 2021. arXiv:2010.14529 , doi:10.1103/PhysRevD.103.122002 .

B. P. Abbott et al. A gravitational-wave standard siren measurement of the Hubble constant. Nature , 551(7678):85-88, 2017. arXiv:1710.05835 , doi:10.1038/nature24471 .

J. 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/ PhysRevD.91.042003 .

Gregory Ashton et al. BILBY: A user-friendly Bayesian inference library for gravitationalwave astronomy. Astrophys. J. Suppl. , 241(2):27, 2019. arXiv:1811.02042 , doi:10.3847/ 1538-4365/ab06fc .

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. , 499(3):3295-3319, 2020. arXiv:2006.00714 , doi:10.1093/mnras/ staa2850 .

Joshua S Speagle. dynesty: a dynamic nested sampling package for estimating Bayesian posteriors and evidences. Monthly Notices of the Royal Astronomical Society , 493(3):3132-3158, Feb 2020. URL: http://dx.doi.org/10.1093/mnras/staa278 , arXiv:1904.02180 , doi:10.1093/mnras/staa278 .

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-1092, 1953.

W. K. Hastings. Monte Carlo sampling methods using Markov chains and their applications. Biometrika , 57(1):97-109, 04 1970. arXiv:https://academic.oup.com/biomet/ article-pdf/57/1/97/23940249/57-1-97.pdf , doi:10.1093/biomet/57.1.97 .

John Skilling. Nested sampling for general Bayesian computation. Bayesian Analysis , 1(4):833 - 859, 2006. doi:10.1214/06-BA127 .

Elena Cuoco, Jade Powell, Marco Cavaglià, Kendall Ackley, Michał Bejger, Chayan Chatterjee, Michael Coughlin, Scott Coughlin, Paul Easter, Reed Essick, et al. Enhancing gravitationalwave science with machine learning. Machine Learning: Science and Technology , 2(1):011002, 5 2020. arXiv:2005.03745 , doi:10.1088/2632-2153/abb93a .

Hunter Gabbard, Chris Messenger, Ik Siong Heng, Francesco Tonolini, and Roderick Murray-Smith. Bayesian parameter estimation using conditional variational autoencoders for gravitational-wave astronomy. Nature Phys. , 18(1):112-117, 2022. arXiv:1909.06296 , doi:10.1038/s41567-021-01425-7 .

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:2002.07656 , doi:10.1103/PhysRevD.102.104057 .

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

Stephen R. Green and Jonathan Gair. Complete parameter inference for GW150914 using deep learning. Mach. Learn. Sci. Tech. , 2(3):03LT01, 2021. arXiv:2008.03312 , doi:10.1088/ 2632-2153/abfaed .

Maximilian Dax, Stephen R. Green, Jonathan Gair, Michael Deistler, Bernhard Schölkopf, and Jakob H. Macke. Group equivariant neural posterior estimation. In International Conference on Learning Representations , 11 2022. arXiv:2111.13139 .

Chayan Chatterjee, Linqing 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. arXiv:2207.14522 .

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 .

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 .

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 .

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/PhysRevD.107.084046 .

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 .

Mark Hannam, Patricia Schmidt, Alejandro Bohé, Leïla Haegel, Sascha Husa, Frank Ohme, Geraint Pratten, and Michael Pürrer. Simple model of complete precessing black-holebinary 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 .

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 .

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 .

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 .

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 ( θ ) that restricts the resulting continuous flow to a one dimensional Gaussian with mean ˆ µ variance ˆ σ 2 . With the optimal transport path µ t ( θ ) = tθ 1 , σ t ( θ ) = 1 -(1 -σ min ) t ≡ σ t from [16], the sample-conditional probability path (5) reads

We set our target distribution

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 function is defined by f ∗ g ( x ) = ∫ f ( x -y ) g ( y ) d y . We define the function

Then the following holds

Marginal probability paths

Marginalizing over θ 1 in (10) with (11), we find

where we defined θ t 1 = tθ 1 and used (13).

Marginal vector field

We now calculate the marginalized vector field u t ( θ ) based on equation (8) in [16]. Using the sample-conditional vector field (6) and the distributions (10) and (11) we find

where we used the change of variables θ ′ 1 = tθ 1 and θ ′′ 1 = θ -θ ′ 1 . Now we evaluate this expression using (12), then the identities (13) and (14) and the marginal probability (15)

By choosing a vector field v t of the form (17) with learnable parameters ˆ µ, ˆ σ 2 , we can thus define a continuous flow that is restricted to 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 1. We first recall the notation from the main part. We always assume that the data is distributed according to p 1 ( θ ) . In addition, there is a known and simple base distribution p 0 and we assume that there is a vector field u t : [0 , 1] × R d → R d that connects p 0 and p 1 in the following sense. We denote by ϕ t the flow generated by u t , i.e., ϕ t satisfies

Then we assume that ( ϕ 1 ) ∗ p 0 = p 1 and we also define the interpolations p t = ( ϕ t ) ∗ p 0 .

We do not have access to the ground truth distributions p t and the vector field u t but we try to learn a vector field v t approximating u t . We denote its flow by ψ t and we define q t = ( ψ t ) ∗ q 0 and q 0 = p 0 . We are interested in the mass covering properties of the learned approximation q 1 of p 1 . In particular, we want to relate the KL-divergence D KL ( p 1 || q 1 ) to the mean squared error,

of the generating vector fields. The first observation is that without any regularity assumptions on v t it is impossible to obtain any bound on the KL-divergence in terms of the mean squared error.

Lemma 1. For every ε > 0 there are vector field u t and v t and a base distribution p 0 = q 0 such that

In addition we can construct u t and v t such that the support of p 1 is larger than the support of q 1 .

Proof. We consider the uniform distribution p 0 = q 0 ∼ U ([ -1 , 1]) and the vector fields and

As before, let ϕ t denote the flow of the vector field u t and similarly ψ t denote the flow of v t . Clearly ϕ t ( θ ) = θ . On the other hand

In particular

This implies that p 1 = ( ϕ 1 ) ∗ p 0 ∼ U ([ -1 , 1]) . On the other hand q 1 = ( ψ 1 ) ∗ q 0 has support in [ -1 , 0] ∪ [ ε, 1] . In particular, the distribution of q 1 is not mass covering with respect to p 1 and D KL ( p 1 || q 1 ) = ∞ . Finally, we observe that the MSE can be arbitrarily small

Here we used that the density of p t (d θ ) is 1 / 2 for -1 ≤ θ ≤ 1 .

We see that an arbitrary small MSE-loss cannot ensure that the probability distribution is mass covering and the KL-divergence is finite. On a high level this can be explained by the fact that for vector fields v t that are not Lipschitz continuous the flow is not necessarily continuous, and we can generate holes in the distribution. Note that we chose p 0 to be a uniform distribution for simplicity, but the result extends to any smooth distribution, in particular the result does not rely on the discontinuity of p 0 .

Next, we investigate the mass covering property for Lipschitz continuous flows. When the flows u t and v t are Lipschitz continuous (in θ ) this ensures that the flows ψ 1 and ϕ 1 are continuous in x and it is not possible to create holes in the distribution as shown above for non-continuous vector fields. We show a weaker bound in this setting.

Lemma 2. For every 0 ≤ δ ≤ 1 there is a base distribution p 0 = q 0 and the are Lipschitz-continuous vector fields u t and v t such that MSE p ( u, v ) = δ and

Proof. We consider p 0 , q 0 and u t as in Lemma 1, and we define

Then we can calculate for 0 ≤ θ ≤ e -2 ε that

Similarly we obtain for ε ≤ θ ≤ 2 ε (solving the ODE f ′ = 2 f )

We find

Next we find for the densities of q 1 that

Together with (30) this implies that the density of q 1 is given by

Note that p 1 ( θ ) = 1 / 2 for -1 ≤ θ ≤ 1 and therefore

and

Moreover we note

which implies (by positivity of the KL-divergence) that

We infer using also p 1 ( θ ) = q 1 ( θ ) = 1 / 2 for θ ∈ [ -1 , 0] ∩ [2 ε, 1] that

On the other hand we can bound

We conclude that

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 KL-divergence by the mean squared error and our Lemmas 1 and 2 are concrete examples. On the other hand, when considering SDEs the KL-divergence 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 u t , v t , and p 0 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 now state our assumptions. We denote the gradient with respect to θ by ∇ = ∇ µ and second derivatives by ∇ 2 = ∇ 2 µν . When applying the chain rule, we leave the indices implicit. We denote by

| · | the Frobenius norm | A | = ( ∑ ij A 2 ij ) 1 / 2 of a matrix. The Frobenius norm is submultiplicative, i.e., | AB | ≤ | A | · | B | and directly generalizes to higher order tensors. Assumption 1 . We assume that

We require one further assumption on p 0 .

Assumption 2 . There is a constant C 1 such that

We also assume that

Note that (41) holds, e.g., if p 0 follows a Gaussian distribution but also for smooth distribution with slower decay at ∞ . If we assume that |∇ ln p 0 ( θ ) | is bounded the proof below simplifies slightly. This is, e.g., the case if p 0 ( θ ) ∼ e -| θ | as | θ |→∞ .

We need some additional notation. It is convenient to introduce ϕ s t = ϕ t ◦ ( ϕ s ) -1 , i.e., the flow from time s to t (in particular ϕ 0 t = ϕ t ) and similarly for ψ . We can now restate and prove Theorem 1.

Theorem 2. Let p 0 = q 0 and assume u t and v t are two vector fields whose flows satisfy p 1 = ( ϕ 1 ) ∗ p 0 and q 1 = ( ψ 1 ) ∗ q 0 . Assume that p 0 satisfies Assumption 2 and u t and v t satisfy Assumption 1. Then there is a constant C > 0 depending on L , L ′ , C 1 , C 2 , and d such that (for MSE p ( u, v ) < 1) )

Remark 1 . We do not claim that our results are optimal, it 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 is not true without the assumption on the second derivative of v t and u t .

Proof. We want to control D KL ( p 1 || q 1 ) . It can be shown that (see equation above (25) in [43] or Lemma 2.19 in [42] )

Using Cauchy-Schwarz we can bound this by

We use the relation (see (3))

which can be equivalently rewritten (setting θ = ϕ t θ 0 ) as

We use the following relation for ∇ ϕ t s

This relation is standard and can be directly deduced from the following ODE for ∇ ϕ t s

We can conclude that for 0 ≤ s, t ≤ 1 the bound

holds. We find

and a similar bound holds for q t . In words, we have shown that the score of p t at θ can be bounded by the score of p 0 of theta transported along the vector field u t minus a correction which quantifies the change of score along the path. We now bound using the definition p t = ( ϕ t ) ∗ p 0 and the assumption (41)

Similarly we obtain using q 0 = p 0

In words, to control the score of q integrated with respect to p t we need to control the distortion we obtain when moving forward with u and backwards with v . We investigate ψ t 0 ϕ t ( θ 0 ) . We now show

First, by definition of ϕ , we find

To evaluate the second contribution we observe

Now (54) follows from (55) and (56) together with ϕ t t = ψ t t = Id . Using (54) we find

Using (50) we conclude that

We use this and the assumption (41) to continue to estimate (53) as follows

Here we used ( a + b + c ) 2 ≤ 3( a 2 + b 2 + c 2 ) in the second to last step. We bound the remaining integral using Cauchy-Schwarz as follows

The last displays together imply

Now we have all the necessary ingredients to bound the derivative of the KL-divergence. We control the second integral in (45) using (51) (and again ( ∑ 4 i =1 a i ) 2 ≤ 4 ∑ a 2 i ) as follows,

Using (52) and (61) we finally obtain

for some constant C > 0 . Finally, we obtain

C SBI Benchmark

In this section, we collect missing details and additional results for the analysis of the SBI benchmark in Section 4.

C.1 Network architecture and hyperparameters

For each task and simulation budget in the benchmark, we perform a mild hyperparameter optimization. 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. 4 the improvement shown by FMPE in MMD is more substantial.

Fig. 7 compares the FMPE results with the optimal transport path from the main text with a comparable score matching model using the Variance Preserving diffusion path [15]. The score matching results were obtained using the same batch size, network size and learning rate as the FMPE network, while optimizing for β min ∈ { 0 . 1 , 1 , 4 } and β max ∈ { 4 , 7 , 10 } . FMPE with the optimal transport path clearly outperforms the score-based model on almost all configurations.

In Fig. 8 we compare FMPE using the architecture proposed in Section 3.2 with ( t, θ ) -conditioning via gated linear units to FMPE with a naive architecture operating directly on the concatenated ( t, θ, x ) vector. For the two displayed tasks the context dimension dim( x ) = 100 is much larger than the parameter dimension dim( θ ) ∈ { 5 , 10 } , and there is a clear performance gain in using the GLU conditioning. Our interpretation is that the low dimensionality of ( t, θ ) means that it is not well-learned by the network when simply concatenated with x .

0

.

00

0

.

05

0

.

10

0

.

15

0

.

20

GL

GL-U

GM

Two Moons

SLCP

10

3

10

4

10

5

0

.

0

0

.

2

0

.

4

0

.

6

B-GLM

10

3

10

4

10

5

B-GLM-Raw

10

3

10

4

10

5

SLCP-D

10

3

10

4

10

5

SIR

10

3

10

4

10

5

LV

Number of Simulations

MMD

NPE

FMPE

Fig. 9 displays the densities of the reference samples under the FMPE model as a histogram for all tasks (extended version of Fig. 3). The support of the learned model q ( θ | x ) covers the reference samples θ ∼ p ( θ | x ) , providing additional empirical evidence for the mass-covering behavior theoretically explored in Thm. 1. However, samples from the true posterior distribution may have a small density under the learned model, especially if the deviation between model and reference is high; see Lotka-Volterra (bottom right panel). Fig. 10 displays P-P plots for two selected tasks.

Finally, we study the impact of our time prior re-weighting for one example task in Fig. 11. We clearly see that our proposed re-weighting leads to increased performance by up-weighting samples for t closer to 1 during training.

0

.

6

0

.

8

1

.

0

GL

0

.

6

0

.

8

1

.

0

GL-U

0

.

6

0

.

8

1

.

0

GM

0

.

6

0

.

8

1

.

0

Two Moons

0

.

6

0

.

8

1

.

0

SLCP

10

3

10

4

10

5

0

.

6

0

.

8

1

.

0

B-GLM

10

3

10

4

10

5

0

.

6

0

.

8

1

.

0

B-GLM-Raw

10

3

10

4

10

5

0

.

6

0

.

8

1

.

0

SLCP-D

10

3

10

4

10

5

0

.

6

0

.

8

1

.

0

SIR

10

3

10

4

10

5

0

.

6

0

.

8

1

.

0

LV

Number of Simulations

C2ST

SMPE

FMPE

10

3

10

4

10

5

0

.

5

0

.

6

0

.

7

0

.

8

0

.

9

1

.

0

B-GLM-Raw

10

3

10

4

10

5

SLCP-D

Number of Simulations

C2ST

GLU embedding

Concatenation

-

20

-

15

-

10

-

5

GL

-

20

-

15

-

10

-

5

-

20

-

15

-

10

-

5

-

15

-

10

-

5

GL-U

-

15

-

10

-

5

-

15

-

10

-

5

0

-

6

-

4

-

2

0

2

GM

-

5

.

0

-

2

.

5

0

.

0

2

.

5

5

.

0

-

5

0

5

-

5

.

0

-

2

.

5

0

.

0

2

.

5

Two Moons

-

4

-

2

0

2

4

-

2

0

2

4

-

9

-

8

-

7

-

6

SLCP

-

10

.

0

-

7

.

5

-

5

.

0

-

2

.

5

-

10

-

5

0

-

20

-

10

0

B-GLM

-

10

0

-

10

-

5

0

5

-

15

-

10

-

5

0

B-GLM-Raw

-

10

-

5

0

5

-

10

-

5

0

5

-

10

-

8

-

6

SLCP-D

-

10

.

0

-

7

.

5

-

5

.

0

-

2

.

5

-

10

-

5

0

-

7

.

5

-

5

.

0

-

2

.

5

0

.

0

SIR

-

6

-

4

-

2

0

2

-

6

-

4

-

2

0

2

-

10

-

5

0

5

LV

-

10

-

5

0

5

log

q

(

θ

x

)

-

5

.

0

-

2

.

5

0

.

0

2

.

5

5

.

0

θ

∼

p

(

θ

|

x

)

θ

∼

q

(

θ

|

x

)

|

0

1

p

0

.

0

0

.

5

1

.

0

CDF(p)

0

1

p

0

.

0

0

.

5

1

.

0

CDF(p)

0

1

p

0

.

0

0

.

5

1

.

0

CDF(p)

0

1

p

0

.

0

0

.

5

1

.

0

CDF(p)

0

1

p

0

.

0

0

.

5

1

.

0

CDF(p)

0

1

p

0

.

0

0

.

5

1

.

0

CDF(p)

10

3

10

4

10

5

0

.

5

0

.

6

0

.

7

0

.

8

0

.

9

1

.

0

Two Moons

Number of Simulations

C2ST

power law

uniform

D Gravitational-wave inference

Wehere 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 of the network is simply a vector field. This also holds for NPSE (model also defined by a vector) and NRE (defined by a scalar). Our FMPE architecture builds on the embedding network developed in [7], however we extend the network capacity by adding more residual blocks (Tab. 3, top panel). For the ( t, θ ) -conditioning 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, θ ) 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. 3, bottom panel). For this configuration, we additionally study the effect of conditioning on ( t, θ ) starting from different layers of the main residual network.

D.2 Data settings

We use the data settings described in [7], with a few minor modifications. In particular, we use the waveform model IMRPhenomPv2 [76-78] and the prior displayed in Tab. 4. Generation of the training dataset with 5,000,000 samples takes around 1 hour on 64 CPUs. Compared to [7], we 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. 5 displays the inference times for FMPE and NPE. NPE requires only a single network pass to produce samples and (log-)probabilities, whereas many forwards passes are needed for FMPE to solve the ODE with a 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, see Eq. (3).

Fig. 12 presents a comparison of the FMPE performance using networks of the same hidden dimensions as the NPE embedding network (Tab. 3 bottom panel). This comparison includes an ablation study on the timing of the ( t, θ ) GLU-conditioning. In the top-row network, the ( t, θ ) conditioning is applied only after the 256-dimensional blocks. In contrast, the middle-row network receives ( t, θ ) immediately after the initial residual block. With FMPE we can achieve performance comparable to NPE, while having only ≈ 1 / 3 of the network size (most of the NPE network parameters are in the flow). This suggests that parameterizing the target distribution in terms of a vector field requires less learning capacity, compared to directly learning its density. Delaying the ( t, θ ) conditioning until the final layers impairs performance. However, the number of FLOPs at inference is considerably reduced, as the context embedding can be cached and a network pass only involves the few layers with the ( t, θ ) conditioning. Consequently, there's a trade-off between accuracy and inference speed, which we will explore in a greater scope in future work.

m

1

m

2

a

1

a

2

t

1

t

2

φ

12

d

L

t

c

α

δ

ψ

FMPE late GLU

FMPE early GLU

NPE

45.3

41.5

1.2

0.5

18.0

11.8

0.3

12.3

8.1

11.4

7.5

0.4

4.8

4.6

0.8

1.0

0.9

0.2

0.3

1.7

2.0

4.8

13.0

0.4

1.2

2.5

3.2

1.6

0.8

0.4

0.3

4.4

9.1

10.1

8.6

0.6

10

20

JSD [mnat]

0

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- Well written paper\n- Nice theoretical results which provide good intuitions despite constrained hypotheses\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nNot applicable.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-03"}

comments: SUMMARY:
This paper proposes solving Bayesian inference problems where we would like to estimate & sample from p(theta | x) when p(theta) and p(x | theta) are known. This is done by randomly sampling from the joint distribution and training conditional generative models as amortized approximate posteriors.

Unlike ELBO-based training methods, the proposed approach does not ever sample from the model during training.

Experiments on some benchmark problems seem to be done, though this reviewer has no experience in these problems.

STRENGTHS:
- Well-written, motivated, and easy to understand.

  - Method is straightforward but seems to work well.

  - Theorem 1 is great.

WEAKNESSES:
- Benchmarks do not seem to be be challenging.

QUESTIONS:
- How is JSD computed? This would typically require sampling from the true posterior. 

  - Are the true posteriors known for these problems?

  - The inference problems may not be challenging enough. Typically, to fit a approximate posterior, we would use the KL(p_model || p_true_posterior) or equivalently, maximizing the ELBO for some given x. This is done because this allows us to sample from the approximate posterior, which has much better properties than sampling from the prior p(theta). I wonder if the reason this is not discussed is because the problems are just not very challenging. For instance, can the proposed approach be used to train a Bayesian neural network?

  - In regards to Theorem 1, have you tried estimating how large the value of C is for the experimental problems? This grows exponentially with respect to time right?

  - In regards to Theorem 1, in the author's opinions, what is the most problematic assumption that may not hold in practice?

LIMITATIONS:
I do not see a discussion on limitations.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nOverall 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.\n\nWEAKNESSES:\nFrom 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.\n\nQUESTIONS:\nN/A\n\nLIMITATIONS:\nPlease refer to the comments provided in the weaknesses section.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}