Year: 2022

Abstract:
Simulation-based inference with conditional neural density estimators is a powerful approach to solving inverse problems in science. However, these methods typically treat the underlying forward model as a black box, with no way to exploit geometric properties such as equivariances. Equivariances are common in scientific models, however integrating them directly into expressive inference networks (such as normalizing flows) is not straightforward. We here describe an alternative method to incorporate equivariances under joint transformations of parameters and data. Our method—called group equivariant neural posterior estimation (GNPE)—is based on self-consistently standardizing the “pose” of the data while estimating the posterior over parameters. It is architecture-independent, and applies both to exact and approximate equivariances. As a real-world application, we use GNPE for amortized inference of astrophysical binary black hole systems from gravitationalwave observations. We show that GNPE achieves state-of-the-art accuracy while reducing inference times by three orders of magnitude.

heading: 1 INTRODUCTION
text: Bayesian inference provides a means of characterizing a system by comparing models against data. Given a forward model or likelihood p(x|θ) for data x described by parameters θ, and a prior p(θ), the Bayesian posterior is proportional to the product, p(θ|x) ∝ p(x|θ)p(θ). Sampling techniques such as Markov Chain Monte Carlo (MCMC) can be used to build up a posterior distribution provided the likelihood and prior can be evaluated.
For models with intractable or expensive likelihoods (as often arise in scientific applications) simulation-based (or likelihood-free) inference methods offer a powerful alternative (Cranmer et al., 2020). In particular, neural posterior estimation (NPE) (Papamakarios & Murray, 2016) uses expressive conditional density estimators such as normalizing flows (Rezende & Mohamed, 2015; Papamakarios et al., 2021) to build surrogates for the posterior. These are trained using model simulations x ∼ p(x|θ), and allow for rapid sampling for any x ∼ p(x), thereby amortizing training costs across future observations. NPE and other density-estimation methods for simulation-based inference (Gutmann & Corander, 2016; Papamakarios et al., 2019; Hermans et al., 2020) have been reported to be more simulation-efficient (Lueckmann et al., 2021) than classical likelihood-free methods such as Approximate Bayesian Computation (Sisson et al., 2018).
Hanford
Livingston
stan dard
ized
dire ctio
n
Virgo
Training an inference network for any x ∼ p(x) can nevertheless present challenges due to the large number of training samples and powerful networks required. The present study is motivated by the problem of gravitational-wave (GW) data analysis. Here the task is to infer properties of astrophysical black-hole mergers based on GW signals observed at the LIGO and Virgo observatories on Earth. Due to the complexity of signal models, it has previously not been possible to train networks to estimate posteriors to the same accuracy as conventional likelihood-based methods (Veitch et al., 2015; Ashton et al., 2019). The GW posterior, however, is equivariant1 under an overall change in the time of arrival of the data. It is also approximately equivariant under a joint change in the sky position and (by triangulation) individual shifts in the arrival times in each detector (Fig. 1). If we could constrain these parameters a priori, we could therefore apply time shifts to align the detector data and simplify the inference task for the remaining parameters.
More generally, we consider forward models with known equivariances under group transformations applied jointly to data and parameters. Our aim is to exploit this knowledge to standardize the pose of the data2 and simplify analysis. The obvious roadblock is that the pose is contained in the set of parameters θ and is therefore unknown prior to inference. Here we describe group equivariant neural posterior estimation (GNPE), a method to self-consistently infer parameters and standardize the pose. The basic approach is to introduce a proxy for the pose—a blurred version—on which one conditions the posterior. The pose of the data is then transformed based on the proxy, placing it in a band about the standard value, and resulting in an easier inference task. Finally, the joint posterior over θ and the pose proxy can be sampled at inference time using Gibbs sampling.
The standard method to incorporating equivariances is to integrate them directly into network architectures, e.g., to use convolutional networks for translational equivariances. Although these approaches can be highly effective, they impose design constraints on network architectures. For GWs, for example, we use specialized embedding networks to extract signal waveforms from frequencydomain data, as well as expressive normalizing flows to estimate the posterior—neither of which is straightforward to make explicitly equivariant. We also have complex equivariance connections between subsets of parameters and data, including approximate equivariances. The GNPE algorithm is extremely general: it is architecture-independent, it applies whether equivariances are exact or approximate, and it allows for arbitrary equivariance relations between parameters and data.
We discuss related work in Sec. 2 and describe the GNPE algorithm in Sec. 3. In Sec. 4 we apply GNPE to a toy example with exact translational equivariance, showing comparable simulation efficiency to NPE combined with a convolutional network. In Sec. 5 we show that standard NPE does not achieve adequate accuracy for GW parameter inference, even with an essentially unlimited number of simulations. In contrast, GNPE achieves highly accurate posteriors at a computational cost three orders of magnitude lower than bespoke MCMC approaches (Veitch et al., 2015). The present paper describes the GNPE method which we developed for GW analysis (Dax et al., 2021), and extends it to general equivariance transformations which makes it applicable to a wide range of problems. A detailed description of GW results is presented in Dax et al. (2021).

heading: 2 RELATED WORK
text: The most common way of integrating equivariances into machine learning algorithms is to use equivariant network architectures (Krizhevsky et al., 2012; Cohen & Welling, 2016). This can be in
1In physics, the term “covariant” is frequently used instead of “equivariant”. 2We adopt the language from computer vision by Jaderberg et al. (2015).
conflict with design considerations such as data representation and flexibility of the architecture, and imposes constraints such as locality. GNPE achieves complete separation of equivariances from these considerations, requiring only the ability to efficiently transform the pose.
Normalizing flows are particularly well suited to NPE, and there has been significant progress in constructing equivariant flows (Boyda et al., 2021). However, these studies consider joint transformations of parameters of the base space and sample space—not joint transformation of data and parameters for conditional flows, as we consider here.
GNPE enables end-to-end equivariances from data to parameters. Consider, by contrast, a conditional normalizing flow with a convolutional embedding network: the equivariance persists through the embedding network but is broken by the flow. Although this may improve learning, it does not enforce an end-to-end equivariance. This contrasts with an invariance, for which the above would be sufficient. Finally, GNPE can also be applied if the equivariance is only approximate.
Several other approaches integrate domain knowledge of the forward model (Baydin et al., 2019; Brehmer et al., 2020) by considering a “gray-box” setting. GNPE allows us to incorporate high-level domain knowledge about approximate equivariances of forward models without requiring access to its implementation or internal states of the simulator. Rather, it can be applied to “black-box” code.
An alternative approach to incorporate geometrical knowledge into classical likelihood-free inference algorithms (e.g., Approximate Bayesian Computation, see (Sisson et al., 2018)) is by constructing (Fearnhead & Prangle, 2012) or learning (Jiang et al., 2017; Chen et al., 2021) equivariant summary statistics s(x), which are used as input to the inference algorithm instead of the raw data x. However, designing equivariant summary statistics (rather than invariant ones) can be challenging, and furthermore inference will be biased if the equivariance only holds approximately.
Past studies using machine-learning techniques for amortized GW parameter inference (Gabbard et al., 2019; Chua & Vallisneri, 2020; Green & Gair, 2021; Delaunoy et al., 2020) all consider simplified problems (e.g., only a subset of parameters, a simplified posterior, or a limited treatment of detector noise). In contrast, the GNPE-based study in Dax et al. (2021) is the only one to treat the full amortized parameter inference problem with accuracy matching standard methods.

heading: 3 METHODS

heading: 3.1 NEURAL POSTERIOR ESTIMATION
text: NPE (Papamakarios & Murray, 2016; Greenberg et al., 2019) is a simulation-based inference method that directly targets the posterior. Given a dataset of prior parameter samples θ(i) ∼ p(θ) and corresponding model simulations x(i) ∼ p(x|θ(i)), it trains a neural density estimator q(θ|x) to estimate p(θ|x). This is achieved by minimizing the loss
LNPE = Ep(θ)Ep(x|θ) [− log q(θ|x)] (1) across the dataset of (θ(i), x(i)) pairs. This maximum likelihood objective leads to recovery of p(θ|x) if q(θ|x) is sufficiently flexible. Normalizing flows (Rezende & Mohamed, 2015; Durkan et al., 2019) are a particularly expressive class of conditional density estimators commonly used for NPE.
NPE amortizes inference: once q(θ|x) is trained, inference is very fast for any observed data xo, so training costs are shared across observations. The approach is also extremely flexible, as it treats the forward model as a black box, relying only on prior samples and model simulations. In many situations, however, these data have known structure that one wants to exploit to improve learning.

heading: 3.2 EQUIVARIANCES UNDER TRANSFORMATION GROUPS
text: In this work we describe a generic method to incorporate equivariances under joint transformations of θ and x into NPE. A typical example arises when inferring the position of an object from image data. In this case, if we spatially translate an image x by some offset ~d—effected by relabeling the pixels— then the inferred position θ should also transform by ~d—by addition to the position coordinates θ. Translations are composable and invertible, and there exists a trivial identity translation, so the set of translations has a natural group structure. Our method works for any continuous transformation group, including rotations, dilations, etc., and in this section we keep the discussion general.
For a transformation group G, we denote the action of g ∈ G on parameters and data as θ → gθ, (2) x→ Tgx. (3)
Here, Tg refers to the group representation under which the data transform (e.g., for image translations, the pixel relabeling). We adopt the natural convention that G is defined by its action on θ, so we do not introduce an explicit representation on parameters. The posterior distribution p(θ|x) is said to be equivariant under G if, when the parameter and data spaces are jointly G-transformed, the posterior is unchanged, i.e., p(θ|x) = p(gθ|Tgx)|det Jg|, ∀g ∈ G. (4) The right-hand side comes from the change-of-variables rule. For translations the Jacobian Jg has unit determinant, but we include it for generality. For NPE, we are concerned with equivariant posteriors, however it is often more natural to think of equivariant forward models (or likelihoods). An equivariant likelihood and an invariant prior together yield an equivariant posterior (App. A.1).
Our goal is to use equivariances to simplify the data—to G-transform x such that θ is taken to a fiducial value. For the image example, this could mean translating the object of interest to the center. In general, θ can also include parameters unchanged under G (e.g., the color of the object), so we denote the corresponding standardized parameters by θ0. These are related to θ by a group transformation denoted gθ, such that gθθ0 = θ. We refer to gθ as the “pose” of θ, and standardizing the pose means to take it to the group identity element e ∈ G. Applying T(gθ)−1 to the data space effectively reduces its dimensionality, making it easier to interpret for a neural network.
Although the preceding discussion applies to equivariances that hold exactly, our method in fact generalizes to approximate equivariances. We say that a posterior is approximately equivariant under G if (4) does not hold, but standardizing the pose nevertheless reduces the effective dimensionality of the dataset. An approximately equivariant posterior can arise if an exact equivariance of the forward model is broken by a non-invariant prior, or if the forward model is itself non-equivariant.

heading: 3.3 GROUP EQUIVARIANT NEURAL POSTERIOR ESTIMATION
text: We are now presented with the basic problem that we resolve in this work: how to simultaneously infer the pose of a signal and use that inferred pose to standardize (or align) the data so as to simplify the analysis. This is a circular problem because one cannot standardize the pose (contained in model parameters θ) without first inferring the pose from the data; and conversely one cannot easily infer the pose without first simplifying the data by standardizing the pose.
Our resolution is to start with a rough estimate of the pose, and iteratively (1) transform the data based on a pose estimate, and (2) estimate the pose based on the transformed data. To do so, we expand the parameter space to include approximate pose parameters ĝ ∈ G. These “pose proxies” are defined using a kernel to blur the true pose, i.e., ĝ = gθ for ∼ κ( ); then p(ĝ|θ) = κ ( (gθ)−1ĝ ) .
The kernel κ( ) is a distribution over group elements, which should be chosen to be concentrated around e; we furthermore choose it to be symmetric. Natural choices for κ( ) include Gaussian and uniform distributions. For translations, the pose proxy is simply the true position with additive noise.
Consider now the posterior distribution p(θ, ĝ|x) over the expanded parameter space. Our iterative algorithm comes from Gibbs sampling this distribution (Roberts & Smith, 1994) (Fig. 2), i.e., alternately sampling θ and ĝ, conditional on the other parameter and x,
θ ∼ p(θ|x, ĝ), (5) ĝ ∼ p(ĝ|x, θ). (6)
The second step just amounts to blurring the pose, since p(ĝ|x, θ) = p(ĝ|θ) = κ ( (gθ)−1ĝ ) . The key first step uses a neural density estimator q that is trained taking advantage of a standardized pose.
For an equivariant posterior, the distribution (5) can be rewritten as (App. A.2) p(θ|x, ĝ) = p ( ĝ−1θ|Tĝ−1x, ĝ−1ĝ ) ∣∣∣det J−1ĝ ∣∣∣ ≡ p(θ′|x′) ∣∣∣det J−1ĝ ∣∣∣ . (7) For the last equality we defined θ′ ≡ ĝ−1θ and x′ ≡ Tĝ−1x, and we dropped the constant argument ĝ−1ĝ = e. This expresses p(θ|x, ĝ) in terms of the ĝ-standardized data x′—which is much easier to
κ = U [−0.2, 0.2]
p(gθ, ĝ|x)
° 2
° 1
0
1
2
g µ
° 2 ° 1 0 1 2
ĝ
°2 °1 0 1 2
gµ
°2 °1
0
1
2
ĝ
°2 °1 0 1 2
gµ
°2 °1
0
1
2
ĝ̂ĝĝg
p(gθ, ĝ|x)
°2 °1 0 1 2
gµ
°2 °1
0
1
2
ĝ
° 2
° 1
0
1
2
g µ
° 2 ° 1 0 1 2
ĝ
κ = U [−1, 1]
°2 °1 0 1 2
gµ
°2 °1
0
1
2
ĝ̂ĝĝg
Figure 2: We infer p(gθ, ĝ|x) with Gibbs sampling by alternately sampling (1) gθ ∼ p(gθ|x, ĝ) (blue) and (2) ĝ ∼ p(ĝ|x, gθ) (orange). For (1) we use a density estimator q(gθ|Tĝ−1(x), ĝ), for (2) the definition ĝ = gθ + , ∼ κ( ). Pose standardization with Tĝ−1 is only allowed due to conditioning on ĝ. Increasing the width of κ accelerates convergence (due to larger steps in parameter space), at the cost of ĝ being a worse approximation for gθ, and therefore pose alignment being less effective.
estimate. We train a neural density estimator q(θ′|x′) to approximate this, by minimizing the loss, LGNPE = Ep(θ)Ep(x|θ)Ep(ĝ|θ) [ − log q ( ĝ−1θ|Tĝ−1x )] . (8) With a trained q(θ′|x′), θ ∼ p(θ|x, ĝ) ⇐⇒ θ = ĝθ′, θ′ ∼ q(θ′|Tĝ−1x). (9)
The estimated posterior is equivariant by construction (App. A.3).
For an approximately-equivariant posterior, (5) cannot be transformed to be independent of ĝ. We are nevertheless able to use the conditioning on ĝ to approximately align x. We therefore train a neural density estimator q(θ|x′, ĝ), by minimizing the loss
LGNPE = Ep(θ)Ep(x|θ)Ep(ĝ|θ) [ − log q ( θ|Tĝ−1x, ĝ )] . (10)
In general, one may have a combination of exact and approximate equivariances (see, e.g., Sec. 5).

heading: 3.4 GIBBS CONVERGENCE
text: The Gibbs-sampling procedure constructs a Markov chain with equilibrium distribution p(θ, ĝ|x). For convergence, the chain must be transient, aperiodic and irreducible (Roberts & Smith, 1994; Gelman et al., 2013). For sensible choices of κ( ) the chain is transient and aperiodic by construction. Further, irreducibility means that the entire posterior can be reached starting from any point, which should be possible even for disconnected posteriors provided the kernel is sufficiently broad. In general, burn-in truncation and thinning of the chain is required to ensure (approximately) independent samples. By marginalizing over ĝ (i.e., ignoring it) we obtain samples from the posterior p(θ|x), as desired.3
Convergence of the chain also informs our choice of κ( ). For wide κ( ), only a few Gibbs iterations are needed to traverse the joint posterior p(θ, ĝ|x), whereas for narrow κ( ) many steps are required (Fig. 2). In the limiting case of κ( ) a delta distribution (i.e., no blurring) the chain does not deviate from its initial position and therefore fails to converge.4 Conversely, a narrower κ( ) better constrains the pose, which improves the accuracy of the density estimator. The width of κ should be chosen based on this practical trade-off between speed and accuracy; the standard deviation of a typical pose posterior is usually a good starting point.
In practice, we obtain N samples in parallel by constructing an ensemble of N Markov chains. We initialize these using samples from a second neural density estimator qinit(gθ|x), trained using standard NPE. Gibbs sampling yields a sequence of sample sets {θ(i)j }Ni=1, j = 0, 1, 2, . . ., each of which represents a distribution Qj(θ|x) over parameters. Assuming a perfectly trained network, one iteration applied to sample set j yields an updated distribution,
Qj+1(θ|x) = p(θ|x) [ Qj(·|x) ∗̄κ p(·|x) ∗̄κ ∗ κ ] (gθ). (11)
The “∗” symbol denotes group convolution and “ ∗̄ ” the combination of marginalization and group convolution (see App. A.4 for details). The true posterior p(θ|x) is clearly a fixed point of this sequence, with the number of iterations to convergence determined by κ and the accuracy of the initialization network qinit.
3In practice, this results only in approximate samples due to the asymptotic behaviour of Gibbs sampling, and a potential mismatch between the trained q and the targeted true posterior.
4This also explains why introducing the pose proxy is needed at all: GNPE would not work without it!

heading: 4 TOY EXAMPLE: DAMPED HARMONIC OSCILLATOR
text: We now apply GNPE to invert a simple model of a damped harmonic oscillator. The forward model gives the time-dependent position x of the oscillator, conditional on its real frequency ω0, damping ratio β, and time of excitation τ . The time series x is therefore a damped sinusoid starting at τ (and zero before). Noise is introduced via a normally-distributed perturbation of the parameters θ = (ω0, β, τ), resulting in a Gaussian posterior p(θ|x) (further details in App. C.1). The model is constructed such that the posterior is equivariant under translations in τ ,
p(ω0, β, τ + ∆τ |T∆τx) = p(ω0, β, τ |x), (12) so we take τ to be the pose. The equivariance of this model is exact, but it could easily be made approximate by, e.g., introducing τ -dependent noise. The prior p(τ) extends from −5 s to 0 s, so for NPE, the density estimator must learn to interpret data from oscillators excited throughout this range.
For GNPE, we shift the data to align the pose near τ = 0 using a Gaussian kernel κ = N [0, (0.1 s)2] (Fig. 3a). We then train a neural density estimator q(θ′|x′) to approximate p(θ′|x′), where θ′ ≡ (ω0, β,− ) and x′ ≡ T−(τ+ )x are pose-standardized. We take q(θ′|x′) to be diagonal Gaussian, matching the known form of the posterior. For each experiment, we train until the validation loss stops decreasing. We also train a neural density estimator qinit(τ |x) with standard NPE on the same dataset to generate initial GNPE seeds. To generate N posterior samples we proceed as follows:
1. Sample τ (i) ∼ qinit(τ |x), i = 1, . . . , N ; 2. Sample (i) ∼ κ( ), set τ̂ (i) = τ (i) + (i), and time-translate the data, x′(i) = T−τ̂(i)x; 3. Sample θ′(i) ∼ q(θ′|x′(i)), and undo the time translation τ̂ (i) to obtain θ(i).
We repeat steps 2 and 3 until the distribution over τ converges. For this toy example and our choice of κ only one iteration is required. For further details of the implementation see App. C.2.
We evaluate GNPE on five simulations by comparing inferred samples to ground-truth posteriors using the c2st score (Friedman, 2004; Lopez-Paz & Oquab, 2017). This corresponds to the test accuracy of a classifier trained to discriminate samples from the target and inferred distributions, and ranges from 0.5 (best) to 1.0 (worst). As baselines we evaluate standard NPE (i) with a network architecture identical to GNPE and (ii) with a convolutional embedding network (NPE-CNN; see App. C.2). Both approaches that leverage the equivariance, GNPE (by standardizing the pose) and NPE-CNN (by using a translation-equivariant embedding network), perform similarly well and far outperform standard NPE (Fig. 3b). This underscores the importance of equivariance awareness. The fact that the NPE network is trained to interpret signals from oscillators excited at arbitrary τ , whereas GNPE focuses on signals starting around τ = 0 (up to a small perturbation) is also reflected in in simplified filters in the first layer of the network (Fig. 3c) and a reduced effective dimension of the input data to GNPE (Fig. 3d).

heading: 5 GRAVITATIONAL-WAVE PARAMETER INFERENCE
text: Gravitational waves—propagating ripples of space and time—were first detected in 2015, from the inspiral, merger, and ringdown of a pair of black holes (Abbott et al., 2016). Since that time, the two LIGO detectors (Hanford and Livingston) (Aasi et al., 2015) as well as the Virgo detector (Acernese et al., 2015) have observed signals from over 50 coalescences of compact binaries involving either black holes or neutron stars (Abbott et al., 2019; 2021a;d). Key scientific results from these observations have included measurements of the properties of stellar-origin black holes that have provided new insights into their origin and evolution (Abbott et al., 2021b); an independent measurement of the local expansion rate of the Universe, the Hubble constant (Abbott et al., 2017); and new constraints on the properties of gravity and matter under extreme conditions (Abbott et al., 2018; 2021c).
Quasicircular binary black hole (BBH) systems are characterized by 15 parameters θ, including the component masses and spins, as well as the space-time position and orientation of the system (Tab. D.1). Given these parameters, Einstein’s theory of general relativity predicts the motion and emitted gravitational radiation of the binary. The GWs propagate across billions of light-years to Earth, where they produce a time-series signal hI(θ) in each of the LIGO/Virgo interferometers I = H,L,V. The signals on Earth are very weak and embedded in detector noise nI . In part to have a tractable likelihood, the noise is approximated as additive and stationary Gaussian. The signal and noise models give rise to a likelihood p(x|θ) for observed data x = {hI(θ) + nI}I=H,L,V. Once the LIGO/Virgo detection pipelines are triggered, classical stochastic samplers are typically employed to determine the parameters of the progenitor system using Bayesian inference (Veitch et al., 2015; Ashton et al., 2019). However, these methods require millions of likelihood evaluations (and hence expensive waveform simulations) for each event analyzed. Even using fast waveform models, it can take O(day) to analyze a single BBH. Faster inference methods are therefore highly desirable to cope with growing event rates, more realistic (and expensive) waveform models, and to make rapid localization predictions for possible multimessenger counterparts. Rapid amortized methods such as NPE have the potential to transform GW data analysis. However, due to the complexity and high dimensionality5 of GW data, it has been a challenge (Gabbard et al., 2019; Chua & Vallisneri, 2020; Green & Gair, 2021; Delaunoy et al., 2020) to obtain results of comparable accuracy and completeness to classical samplers. We now show how GNPE can be used to exploit equivariances to greatly simplify the inference problem and achieve for the first time performance indistinguishable from “ground truth” stochastic samplers—at drastically reduced inference times.

heading: 5.1 EQUIVARIANCES OF SKY POSITION AND COALESCENCE TIME
text: We consider the analysis of BBH systems. Included among the parameters θ are the time of coalescence tc (as measured at geocenter) and the sky position (right ascension α, declination δ). Since GWs propagate at the speed of light, these are related to the times of arrival tI of the signal in each of the interferometers.6 Our priors (based on the precision of detection pipelines) constrain tI to a range of ≈ 20 ms, which is much wider than typical posteriors. Standard NPE inference networks must therefore be trained on simulations with substantial time shifts.
The detector coalescence times tI—equivalently, (tc, α, δ)—can alternatively be interpreted as the pose of the data, and standardized using GNPE. The group G transforming the pose factorizes into a direct product of absolute and relative time shifts,
G = Gabs ×Grel. (13) Group elements gabs ∈ Gabs act by uniform translation of all tI , whereas grel ∈ Grel act by individual translation of tL and tV. We work with data in frequency domain, where time translations act by multiplication, TgxI = e−2πif∆tIxI . Absolute time shifts correspond to a shift in tc, and are an exact equivariance of the posterior, p(gabsθ|Tgabsx) = p(θ|x). Relative time shifts correspond to a change in (α, δ) (as well as tc). This is only an approximate equivariance, since a change in sky position changes the projection of the incident signal onto the detector arms, leading to a subdominant change to the signal morphology in each detector.
5In our work, we analyze 8 s data segments between 20 Hz and 1024 Hz. Including also noise information, this results in 24,099 input dimensions per detector.
6We consider observations made in either nI = 2 or 3 interferometers.
NPE NPE (chained)
GNPE (wide κ)
GNPE (narrow κ)
0.50
0.55
0.60
0.65
0.70
c2 st
∅ 0.627
∅ 0.558
∅ 0.526 ∅ 0.519
GW150914 GW151012 GW170104 GW170729
GW170809 GW170814 GW170818 GW170823
Figure 4: Comparison of estimated posteriors against LALINFERENCE MCMC for eight GW events, as quantified by c2st (best: 0.50, worst: 1.00). GNPE with a wide kernel outperforms both NPE baselines, while being only marginally slower (1 iteration ∼ 2 s). With a narrow kernel and 30 iterations (∼ 60 s), we achieve c2st< 0.55 across all events. ∅ indicates the average across all eight events. For an alternative metric (MSE) see Fig. D.2.

heading: 5.2 APPLICATION OF GNPE
text: We use GNPE to standardize the pose within a band around tI = 0. We consider two modes defined by different uniform blurring kernels. The “accurate” mode uses a narrow kernel κnarrow = U [−1 ms, 1 ms]nI , whereas the “fast” mode uses a wide kernel κwide = U [−3 ms, 3 ms]nI . The latter is intended to converge in just one GNPE iteration, at the cost of having to interpret a wider range of data.
We define the blurred pose proxy ĝI ≡ tI + I , where I ∼ κ( I). We then train a conditional density estimator q(θ′|x′, ĝrel), where θ′ = ĝ−1abs θ and x′ = Tĝ−1x. That is, we condition q on the relative time shift (since this is an approximate equivariance) and we translate parameters by the absolute time shift (since this is an exact equivariance). We always transform the data by the full time shift. We train a density estimator qinit({tI}I=H,L,V|x) using standard NPE to infer initial pose estimates. The difficulty of the inference problem (high data dimensionality, significant noise levels, complex forward model) combined with high accuracy requirements to be scientifically useful requires careful design decisions. In particular, we initialize the first layer of the embedding network with principal components of clean waveforms to provide an inductive bias to extract useful information. We further use an expressive neural-spline normalizing flow (Durkan et al., 2019) to model the complicated GW posterior structure. See App. D.2 for details of network architecture and training.

heading: 5.3 RESULTS
text: We evaluate performance on all eight BBH events from the first Gravitational-Wave Transient Catalog (Abbott et al., 2019) consistent with our prior (component masses greater than 10 M ). We generate reference posteriors with the LIGO/Virgo MCMC code LALINFERENCE (Veitch et al., 2015). We quantify the deviation between NPE samples and the reference samples using c2st.
We compare performance against two baselines, standard NPE and a modified approach that partially standardizes the pose (“chained NPE”). For the latter, we use the chain rule to decompose the posterior, p(θ|x) = p(φ, λ|x) = p(φ|x, λ) · p(λ|x), (14) where λ = (tc, α, δ) are the pose parameters and φ ⊂ θ collects the remaining 12 parameters. We use standard NPE to train a flow q(λ|x) to estimate p(λ|x), and a flow q(φ|x′, λ) to estimate p(φ|x, λ). The latter flow is conditioned on λ, which we use to standardize the pose of x. In contrast to GNPE, this approach is sensitive to the initial pose estimate q(λ|x), which limits accuracy (Figs. 4 and D.7). We note that all hyperparameters of the flow and training protocol (see App. D.2) were extensively optimized on NPE, and then transferred to GNPE without modification, resulting in conservative estimates of the performance advantage of GNPE. Fast-mode GNPE converges in one iteration, whereas accurate-mode requires 30 (convergence is assessed by the JS divergence between the inferred pose posteriors from two successive iterations).
Standard NPE performs well on some GW events but lacks the required accuracy for most of them, with c2st scores up to 0.71 (Fig. 4). Chained NPE performs better across the dataset, but performs poorly on events such as GW170814, for which the initial pose estimate is inaccurate. Indeed, we find
that inaccuracies of that baseline can be almost entirely attributed to the initial pose estimate (Fig. D.6). Fast-mode GNPE with only one iteration is already more robust to this effect due to the blurring operation of the pose proxy (Fig. D.7). Both GNPE models significantly outperform the baselines, with accurate-mode obtaining c2st scores < 0.55 across all eight events. We emphasize that the c2st score is sensitive to any deviation between reference samples and samples from the inferred posterior. On a recent benchmark by Lueckmann et al. (2021) on examples with much lower parameter and data dimensions, even state-of-the-art SBI algorithms rarely reached c2st scores below 0.6. The fact that GNPE achieves scores around 0.52—i.e., posteriors which are nearly indistinguishable from the reference—on this challenging, high-dimensional, real-world example underscores the power of exploiting equivariances with GNPE.
Finally, we visualize posteriors for two events, GW170809 and GW170814, in Fig. 5. The quantitative agreement between GNPE and MCMC (Fig. 4) is visible from the overlapping marginals for all parameters displayed. NPE, by contrast, deviates significantly from MCMC in terms of shape and position. Note that we show a failure case of NPE here; for other events, such as GW170823, deviations of NPE from the reference posterior are less clearly visible.

heading: 6 CONCLUSIONS
text: We described GNPE, an approach to incorporate exact—and even approximate—equivariances under joint transformations of data and parameters into simulation-based inference. GNPE can be applied to black-box scientific forward models and any inference network architecture. It requires similar training times compared to NPE, while the added complexity at inference time depends on the number of GNPE iterations (adjustable, but typically O(10)). We show with two examples that exploiting equivariances with GNPE can yield large gains in simulation efficiency and accuracy.
For the motivating problem of GW parameter estimation, GNPE achieves for the first time rapid amortized inference with results virtually indistinguishable from MCMC (Dax et al., 2021). This is an extremely challenging “real-world” scientific problem, with high-dimensional input data, complex signals, and significant noise levels. It combines exact and approximate equivariances, and there is no clear path to success without their inclusion along with GW-specialized architectures and expressive density estimators.

heading: ETHICS STATEMENT
text: Our method is primarily targeted at scientific applications, and we do not foresee direct applications which are ethically problematic. In the context of GW analysis, we hope that GNPE contributes to reducing the required amount of compute, in particular when the rate of detections increases with more sensitive detectors in the future.

heading: REPRODUCIBILITY STATEMENT
text: The experimental setup for the toy model is described in App. C.2. We also provide the code at https://tinyurl.com/wmbjajv8. The setup for GW parameter inference is described in App. D.2. We will release a python package including scripts for the experiments carried out in section 5.

heading: ACKNOWLEDGMENTS
text: We thank A. Buonanno, T. Gebhard, J.M. Lückmann, S. Ossokine, M. Pürrer and C. Simpson for helpful discussions. We thank the anonymous reviewer for coming up with the illustration of GNPE in App. B. This research has made use of data, software and/or web tools obtained from the Gravitational Wave Open Science Center (https://www.gw-openscience.org/), a service of LIGO Laboratory, the LIGO Scientific Collaboration and the Virgo Collaboration. LIGO Laboratory and Advanced LIGO are funded by the United States National Science Foundation (NSF) as well as the Science and Technology Facilities Council (STFC) of the United Kingdom, the Max-Planck-Society (MPS), and the State of Niedersachsen/Germany for support of the construction of Advanced LIGO and construction and operation of the GEO600 detector. Additional support for Advanced LIGO was provided by the Australian Research Council. Virgo is funded, through the European Gravitational Observatory (EGO), by the French Centre National de Recherche Scientifique (CNRS), the Italian Istituto Nazionale di Fisica Nucleare (INFN) and the Dutch Nikhef, with contributions by institutions from Belgium, Germany, Greece, Hungary, Ireland, Japan, Monaco, Poland, Portugal, Spain. This material is based upon work supported by NSF’s LIGO Laboratory which is a major facility fully funded by the National Science Foundation. M. Dax thanks the Hector Fellow Academy for support. M. Deistler thanks the International Max Planck Research School for Intelligent Systems (IMPRS-IS) for support. B.S. and J.H.M. are members of the MLCoE, EXC number 2064/1 – Project number 390727645. This work was supported by the German Federal Ministry of Education and Research (BMBF): Tübingen AI Center, FKZ: 01IS18039A. We use PyTorch (Paszke et al., 2019), nflows (Durkan et al., 2020) and sbi (Tejero-Cantero et al., 2020) for the implementation of our neural networks. The plots are generated with matplotlib (Hunter, 2007) and ChainConsumer (Hinton, 2016).

heading: A DERIVATIONS

heading: A.1 EQUIVARIANCE RELATIONS
text: Consider a system with an exact equivariance under a joint transformation of parameters θ and observations x,
θ → θ′ = gθ, (15) x→ x′ = Tgx. (16)
An invariant prior fulfills the relation
p(θ) = p(θ′) |det Jg| , (17) where the Jacobian Jg arises from the change of variables rule for probability distributions. A similar relation holds for an equivariant likelihood,
p(x|θ) = p(x′|θ′) |det JT | . (18) An invariant prior and an equivariant likelihood further imply for the evidence p(x)
p(x) = ∫ p(x|θ)p(θ)dθ = ∫ p(x′|θ′) |det JT | p(θ′) |det Jg| dθ = p(x′) |det JT | . (19)
Combining an invariant prior with an equivariant likelihood thus leads to the equivariance relation
p(θ|x) = p(x|θ)p(θ) p(x) = p(x′|θ′) |det JT | p(θ′) |det Jg|
p(x′) |det JT | = p(θ′|x′) |det Jg|
(20)
for the posterior, where we used Bayes’ theorem and equations (17), (18), and (19).
A.2 EQUIVARIANCE OF p(θ|x, ĝ)
Here we derive that an equivariant posterior p(θ|x) remains equivariant if the distribution is also conditioned on the proxy ĝ, as used in equation (7). With the definition p(ĝ|θ) = κ ( (gθ)−1ĝ ) from section 3.3, p(ĝ|θ) is equivariant under joint application of h ∈ G to ĝ and θ,
p(ĝ|θ) = κ ( (gθ)−1ĝ ) = κ ( (gθ)−1h−1hĝ ) = κ ( (hgθ)−1hĝ ) = p(hĝ|hθ). (21)
where for the last equality we used ghθ = hgθ. This implies, that p(ĝ|x) is equivariant under joint application of h and Th,
p(ĝ|x) = ∫ p(ĝ|θ, x)p(θ|x) dθ (21),(4)= ∫ p(hĝ|hθ)p(hθ|Thx) |det Jh| dθ
= p(hĝ|Thx). (22)
in the second step we used p(ĝ|θ, x) = p(ĝ|θ). From these relations, the equivariance relation used in equation (7) follows,
p(θ|x, ĝ) = p(θ, ĝ|x) p(ĝ|x) = p(ĝ|θ, x)p(θ|x) p(ĝ|x) (21),(4),(22) = p(hĝ|hθ)p(hθ|Thx) p(hĝ|Thx) |det Jh|
= p(hθ|Thx, hĝ) |det Jh| . (23)

heading: A.3 EXACT EQUIVARIANCE OF INFERRED POSTERIOR
text: Consider a posterior that is exactly equivariant under G,
p(θ|x) = p(gθ|Tgx)|det Jg|, ∀g ∈ G. (24) We here show that the posterior estimated using GNPE is equivariant under G by construction. This holds regardless of whether q(θ′|x′) has fully converged to p(θ′|x′).
With GNPE, the equivariant posterior p(θ|xo) for an observation xo is inferred by alternately sampling
θ(i) ∼ p(θ|xo, ĝ(i−1)) ⇐⇒ θ(i) = ĝ(i−1)θ′(i), θ′(i) ∼ q(θ′|T(ĝ(i−1))−1xo), (25) ĝ(i) ∼ p(ĝ|xo, θ(i)) ⇐⇒ ĝ(i) = gθ (i) , ∼ κ( ), (26)
see also equation (9). Now consider a different observation x̃o = Thxo that is obtained by altering the pose of xo with Th, where h is an arbitrary element of the equivariance group G. Applying the joint transformation
ĝ → hĝ, (27) xo → Thxo, (28)
in (25) leaves θ′ invariant,
q(θ′|T(hĝ)−1Thxo) = q(θ′|Tĝ−1(Th)−1Thxo) = q(θ′|Tĝ−1xo). (29)
We thus find that θ in (25) transforms equivariantly under joint application of (27) and (28),
θ = ĝθ′ → (hĝ)θ′ = h(ĝθ′) = hθ. (30)
Conversely, applying
θ → hθ (31)
in (26) transforms ĝ by
ĝ = gθ → g(hθ) = hgθ = hĝ. (32)
The θ samples (obtained by marginalizing over ĝ) thus transform θ → hθ under xo → Thx0, which is consistent with the desired equivariance (24).
Another intuitive way to see this is to consider running an implementation of the Gibbs sampling steps (25) and (26) with fixed random seed for two observations xo (initialized with ĝ(0) = ĝxo) and Thxo (initialized with ĝ(0) = hĝxo). The Gibbs sampler will yield parameters samples (θi) N i=1 for xo, and the exact same samples (hθi)Ni=1 for Thxo, up to the global transformation by h. The reason is that the density estimator q(θ′|x′) is queried with the same x′ for both observations xo and Thxo in each iteration i. Since the truncated, thinned samples are asymptotically independent of the initialization, this shows that (24) is fulfilled by construction.
A.4 ITERATIVE INFERENCE AND CONVERGENCE
GNPE leverages a neural density estimator of the form q(θ|x′, ĝ) to obtain samples from the joint distribution p(θ, ĝ|x). This is done by iterative sampling as described in section 3. Here we derive equation (11), which states how a distribution Qj(θ|x) is updated by a single GNPE iteration. Given a distribution Qθj (θ|x) of θ samples in iteration j, we infer samples for the pose proxy ĝ for the next iteration by (i) extracting the pose gθ from θ (this essentially involves marginalizing over all non pose related parameters) and (ii) blurring the pose gθ with the kernel κ, corresponding to a group convolution.A.1 We denote this combination of marginalization and group convolution with the “ ∗̄ ” symbol,
Qĝj+1(ĝ|x) = ∫ dθ Qθj (θ|x)κ((gθ)−1ĝ) = ( Qθj (·|x) ∗̄κ ) (ĝ). (33)
For a given proxy sample ĝ, a (perfectly trained) neural density estimator infers θ with
p(θ|x, ĝ) = p(θ, ĝ|x) p(ĝ|x) = p(ĝ|x, θ)p(θ|x) p(ĝ|x) = p(θ|x) κ((gθ)−1ĝ) (pθ(·|x) ∗̄κ) (ĝ) , (34)
A.1We define a group convolution as (A ∗B)(ĝ) = ∫ dg A(g)B(g−1ĝ), which is the natural extension of a
standard convolution.
where we used p(ĝ|θ) = κ((gθ)−1ĝ). Combining (33) and (34), the updated distribution over θ samples reads
Qθj+1(θ|x) = ∫ dĝ p(θ|x, ĝ)Qĝj+1(ĝ|x)
= ∫ dĝ ( Qθj (·|x) ∗̄κ ) (ĝ) p(θ|x) κ((g θ)−1ĝ)
((pθ(·|x) ∗̄κ))(ĝ) = p(θ|x) ∫ dĝ ( Qθj (·|x) ∗̄κ ) (ĝ)
((pθ(·|x) ∗̄κ))(ĝ) κ((g θ)−1ĝ)
= p(θ|x) ( Qθj (·|x) ∗̄κ pθ(·|x) ∗̄κ ∗ κ (−) ) (ĝ).
(35)
Here, κ(−) denotes the reflected kernel, κ(−)(g) = κ(g−1)∀g. Since we choose a symmetric kernel in practice, we use κ = κ(−) in (11).
In this notation, the initialization of the pose gθ in iteration 0 with qinit simply means setting Q0(·|x) ∗̄κ = qinit(·|x) ∗ κ.

heading: B GNPE FOR SIMPLE GAUSSIAN LIKELIHOOD AND PRIOR
text: Consider a simple forward model τ → x with a given prior p(τ) = N (−5, 1)[τ ] (36)
and a likelihood
p(x|τ) = N (τ, 1)[x], (37) where the normal distribution is defined by
N (µ, σ2)[x] = exp
( −(x−µ)2
2σ2 ) √
2πσ . (38)
The evidence can be computed from the prior (36) and likelihood (37), and reads
p(x) = ∫ dτp(τ)p(x|τ) = ∫ dτN (−5, 1)[τ ]N (τ, 1)[x] = N (−5, 2)[x]. (39)
The posterior is then given via Bayes’ theorem and reads
p(τ |x) = p(x|τ)p(τ) p(x) = N (τ, 1)[x]N (−5, 1)[τ ] N (−5, √ 2)[x]
= N ( x− 5
2 , 1/2
) [τ ]. (40)

heading: B.1 EQUIVARIANCES
text: The likelihood (37) is equivariant under G, i.e., the joint transformation τ → gτ = τ + ∆τ, x→ T lgx = x+ ∆τ.
(41)
This follows directly from (37) and (38). If the prior was invariant, then this equivariance would be inherited by the posterior, see App. A.1. However, the prior is not invariant. It turns out that the posterior is still equivariant, but x transforms under a different representation than it does for the equivariance of the likelihood. Specifically, the posterior (40) is equivariant under joint transformation
τ → gτ = τ + ∆τ, x→ T pg x = x+ 2 ·∆τ,
(42)
which again directly follows from (40) and (38). Importantly, T pg 6= T lg , i.e., the representation under which x transforms is different for the equivariance of the likelihood and the posterior. For GNPE, the relevant equivariance is that of the posterior, i.e. the set of transformations (42), see also equation (4). The equivariance relation of the posterior thus reads
p(τ |x) = p(gτ |T pg x)|det Jg|, ∀g ∈ G. (43)

heading: B.2 GNPE
text: We choose τ as the pose, which we aim to standardize with GNPE. We define the corresponding proxy as
τ̂ = τ + , ∼ κ( ) = N (0, 1)[ ]. (44) We can use GNPE to incorporate the exact equivariance of the posterior by construction. To that end we define
τ ′ = g(−τ̂)τ = τ − τ̂ , x′ = T p
g(−τ̂) x = x− 2 · τ̂ . (45)
We then train a neural density estimator to estimate p(τ ′|x′). This distribution is of the same form as p(τ |x) and simply given by
p(τ ′|x′) (43),(40)= N ( x′ − 5
2 , 1/2
) [τ ′] (46)
due to the equivariance (43). We here assume a neural density estimator that estimates (46) perfectly. For GNPE, we
1. Initialize τ (1) = 0; 2. Sample τ̂ (1) by τ̂ (1) = τ (1) + , ∼ N (0, 1)[ ], and compute τ ′ and x′ via (45); 3. Sample τ (2) by τ (2) = τ ′(2) + τ̂ (1), with τ ′(2) ∼ p(τ ′|x′) = N ( x′−5 2 , 1/2 ) [τ ′];
and repeat (2) and (3) multiple times. This constructs a Markov chain. To obtain (approximately independent) posterior samples τ ∼ p(τ |x), we truncate to account for burn-in, thin the chain and marginalize over τ̂ . We find that the chain indeed converges to the correct posterior (40), see Fig. B.1.

heading: C TOY EXAMPLE

heading: C.1 FORWARD MODEL
text: The toy model in section 4 describes the motion of a damped harmonic oscillator that is initially at rest and excited at time τ with an infinitely short pulse. The time evolution of that system is governed
by the differential equation
d2 dt2 x(t) + 2βω0 d dt x(t) + ω20x(t) = δ(t− τ), (47)
where ω0 denotes the undamped angular frequency and β the damping ratio. The solution for the time series x(t) is given by the Green’s function for the corresponding differential operator and reads
x(t) = 0, t ≤ τe−βω0(t−τ) · sin(√1−β2ω0(t−τ))√ 1−β2ω0 , t > τ. (48)
This equation describes a deterministic, injective mapping between parameters θ = (ω0, β, τ) and a time series observation x, x = f(θ). (49) This implies a likelihood p(x|θ) = δ(x− f(θ)), and thus a point-like posterior. To showcase (G)NPE on this toy problem we introduce stochasticity by setting
x = f(θ + δθ) (50)
instead. We sample δθ from an uncorrelated Gaussian distribution
δθ ∼ N (0,Σ), Σ =  σ2ω0 0 00 σ2β 0 0 0 σ2τ  , (51) with σω0 = 0.3 Hz, σβ = 0.03 and στ = 0.3 s. Due to the injectivity of f , the posterior p(θ|x) reduces to the probability p(δθ = f−1(x) − θ). With a uniform prior, and neglecting boundary effects, this implies an uncorrelated Gaussian posterior p(θ|x) centered around f−1(x) with standard deviations as specified above. We choose this approach over, e.g., adding noise straight to observations to keep the problem as simple as possible, such that the focus remains on the comparison of GNPE and NPE. In particular, knowing that the ground truth posteriors are Gaussian, we can use a simple Gaussian density estimator.
We choose uniform priors
p(ω0) = U [3, 10] Hz, p(β) = U [0.2, 0.5], p(τ) = U [−5, 0] s. (52) The observational data x is the discretized time series in the interval [−5,+5] s with 2000 evenly sampled bins. When applying time shifts with GNPE, we impose cyclic boundary conditions.
C.2 IMPLEMENTATION
We use a Gaussian density estimator for all methods (since we know that the true posterior is Gaussian). For NPE, we use a feedforward neural network with [128, 32, 16] hidden units and with ReLU activation functions as an embedding network. For NPE-CNN, we use a three-layer convolutional embedding network with kernel sizes [5,5,5], stride 1, [6,12,12] channels, average pooling with kernel size 7 and stride 7, and ReLU activation functions. For GNPE, we use the same architecture as for NPE for both, q(θ′|x′) and qinit(τ |x). For further hyperparameters, we use the defaults of the sbi package (Tejero-Cantero et al., 2020).

heading: C.3 RESULTS
text: For all methods, we compute the average classifier two-sample test score (c2st) based on 10,000 samples from the estimated and the ground truth posterior for five different simulations. We then average the accuracy across 10 different seeds.

heading: D GRAVITATIONAL WAVE PARAMETER INFERENCE

heading: D.1 FORWARD MODEL AND AMORTIZATION
text: The forward model mapping binary black hole parameters θ (Tab. D.1) to simulated measurements x in the detectors consists of two stages. Firstly, the waveform polarizations h(θ) for given parameters
θ are computed with the waveform model IMRPhenomPv2 (Hannam et al., 2014; Khan et al., 2016; Bohé et al., 2016). Secondly, the signals are projected onto the detectors, and noise is added to obtain a realistic signal x. To a good approximation, we assume the noise to be Gaussian and stationary over the duration of a single GW signal. However, the noise spectrum, determined by the power spectral density (PSD) Sn, drifts over the duration of an observing run. To fully amortize the computational cost, we use a variety of different PSDs Sn in training, and additionally condition the inference network on Sn. At inference time, this enables instant tuning of the inference network to the PSD estimated at the time of the event, see Dax et al. (2021) for details. Since this conditioning on Sn has no effect on the GNPE algorithm outlined in this work, we keep it implicit in all equations.

heading: D.2 NETWORK ARCHITECTURE AND TRAINING
text: The inference network consists of an embedding network, that reduces the high dimensional input data to a 128 dimensional feature vector, and the normalizing flow, that takes this feature vector as input. For each detector, the input to the embedding network consists of the complex-valued frequency domain strain in the range [20 Hz, 1024 Hz] with a resolution of 0.125 Hz, and PSD information (1046 · Sn)−1/2 with the same binning. This results to a total of (3 · 8,033) = 24,099 real input bins per detector. The first module of the embedding network consists of a linear layer per detector, that maps this 24,096 dimensional input to 400 components. We initialize this compression layer with PCA components of raw waveforms. This provides a strong inductive bias to the network to filter out GW signals from extremely noisy data. Note that this important step is only possible since GNPE is architecture independent—it is for instance not compatible with a convolutional neural network. Following this compression layer, we use a series of 24 fully-connected residual blocks with two layers each to compress the output to the desired 128 dimensional feature vector. We use batch normalization and ELU activation functions. Importantly, the conditioning of the flow on the proxy ĝrel. is done after the embedding network, by concatenating ĝrel. to the embedded feature vector.
Following this, we use a neural spline flow (Durkan et al., 2019) with rational-quadratic spline coupling transforms as density estimator. We use 30 such transforms, each of which is associated with 5 two-layer residual blocks with hidden dimension 512. In total, the inference network has 348 hidden layers and 1.31 ·108 (for two detectors) or 1.42 ·108 (for three detectors) learnable parameters. We train the inference network with a data set of 5 · 106 waveforms with parameters θ sampled from the priors specified in table D.1, and reserve 2% of the data for validation. We pretrain the network with learning rate of 3 · 10−4 for 300 epochs with fixed PSD, and finetune for another 150 epochs with learning rate of 3 · 10−5 with varying PSDs. With batch size 4,096, training takes 16-18 days on a NVIDIA Tesla V100 GPU.

heading: D.3 RESULTS
text: The c2st scores between inferred posterior and the MCMC reference shown in Fig. 4 are computed using the code and default hyperparameters of Lueckmann et al. (2021). For each event, we compute the c2st score of 10,000 samples for inferred and target posterior. Fig. 4 displays the mean of the score across 5 different sample realizations, Fig. D.1 additionally shows the corresponding standard deviation. For technical reasons we use only 12 of the 15 inferred parameters; specifically we omit the geocentric time of coalescence tc (since the reference posteriors generated with LALInference do not contain that variable) and the sky position parameters α and δ (since the NPE baseline with chain rule decomposition infers these in another basis).

Review summaries:
The authors propose group equivariant neural posterior estimation (GNPE), a posterior estimation method which can self-consistently infer parameters and standardize the pose. For equivariant posterior distributions, the GNPE can achieve better performance than the traditional NPE. Moreover, GNPE can also be applied to cases where the equivariance of the posterior is approximately estimated.  The major advantage of GNPE over the other flow based models under equivariance constraints is that GNPE is architecture-indepenent, i.e., an arbitrary flow model can be utilized in GNPE and we don't need to design a special network structure for the equivariance.

In this paper the authors present a method for performing Bayesian inference in likehood-free settings where the data and parameters are jointly equivariant or approximately equivariant.  Examples include translational or rotational equivariance, where if the latent parameter are translate or rotated, then the distribution over the data is the same, except that the data are translated or rotated in the same way as the latent parameters.  The method presented here acts by mapping each simulation to a standard "pose" and then aims to learn both the original pose and any additional parameters from the data.  There are two main components to this method -- 1) the standard simulation-based inference method of estimating latent parameters from data, but applied to data that have had their pose standardized and 2) a Gibbs sampling framework to use the family of posteriors learned from the first component to jointly estimate the pose and ant other parameters from unposed data.