G ROUP EQUIVARIANT NEURAL POSTERIOR ESTIMATION
Anonymous authors
Paper under double-blind review
A BSTRACT
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 scientiﬁc models, however integrating them directly into expressive inference networks (such as normalizing ﬂows) 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.
I NTRODUCTION
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 model parameters θ , and a prior p ( θ ) over parameters, the Bayesian posterior distribution is proportional to their 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 scientiﬁc 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 ﬂows (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-efﬁcient (Lueckmann et al., 2021) than classical likelihood-free methods such as Approximate Bayesian Computation (Sisson et al., 2018).
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 applying NPE to the problem of gravitational-wave (GW) data analysis.
Here the task is to infer properties of astrophysical black-hole merger events 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 equivariant 1 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.
[Figure]
s t a n d a r d
i z
e d d i r e c ti o n
[Figure]
algorithm
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 data 2 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 ﬂows to estimate the posterior—none 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 section 2 and describe the GNPE algorithm in section 3.
In section 4 we present a toy example with an exact translational equivariance.
We show that GNPE signiﬁcantly outperforms equivariance-agnostic NPE in terms of simulation efﬁciency, and performs on par with
NPE equipped with a convolutional neural network.
In section 5 we show that standard NPE does not achieve adequate accuracy for GW parameter inference, even with an essentially unlimited number of model simulations.
In contrast, by exploiting physical equivariances, 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, and extends it to general equivariance transformations which make it applicable to a wide range of applications.
A detailed description of its results on empirical GW data is presented in the companion physics paper (Anonymous et al., 2021).
We conclude in section 6.
R ELATED WORK
As noted above, 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 conﬂict with other design considerations, such as data representation and ﬂexibility of the architecture, and impose constraints such as locality.
GNPE achieves complete separation of equivariances from these considerations, requiring only the ability to efﬁciently transform the pose.
Normalizing ﬂows are particularly well suited to NPE, and there has been signiﬁcant recent progress in constructing equivariant ﬂows (Boyda et al., 2021).
However, those studies have focused on joint transformation of base-space and sample-space parameters under the ﬂow— not on joint transformation of data and sample-space parameters for a conditional ﬂows, as we consider here.
GNPE provides the ability for end-to-end equivariances—from data to parameters.
For example, when using a convolutional embedding network in conjunction with a conditional normalizing ﬂow, the equivariance persists up until it is passed as context to the ﬂow, which would make it easier for the ﬂow to learn.
However, it does not fully enforce an end-to-end equivariance.
This is in contrast to invariances , for which the above approach would be sufﬁcient.
Finally, and importantly, GNPE can also be applied if the equivariance is only approximate —as demonstrated in our application.
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 simpliﬁed problems (e.g., only a subset of parameters, a simpliﬁed posterior, or a limited treatment of detector noise).
In contrast, the GNPE-based study in Anonymous et al.
(2021) is the only one to treat the full amortized parameter inference problem with accuracy matching standard methods.
M ETHODS 3. N EURAL POSTERIOR ESTIMATION
NPE (Papamakarios & Murray, 2016) 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
[Equation]
across the dataset of ( θ ( i ) , x ( i ) ) pairs.
This maximum likelihood objective leads to recovery of p ( θ | x ) if q ( θ | x ) is sufﬁciently ﬂexible.
Normalizing ﬂows (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 x o , so training costs are shared across all observations.
Furthermore, the approach treats the model as a black box, only samples from the forward model are required.
Although this provides great ﬂexibility, it does not allow one to incorporate knowledge about approximate or exact equivariances.
3. G ROUP TRANSFORMATIONS AND POSE STANDARDIZATION
Consider for concreteness the analysis of image data, and suppose that we wish to infer the position of an object.
The posterior 3 should be equivariant under translations applied jointly to the image itself and the inferred position.
Namely, if we relabel the pixels of the image to apply an overall translation, then the posterior probability of the position is shifted by the same amount.
Translations are composable and invertible, and there exists a trivial identity translation, so the set of translations has a natural group structure.
More generally, one may be interested in other continuous transformation groups, such as rotations, scalings, etc., so in this work we will consider generic Lie groups.
For a transformation group G , we denote the action of g ∈ G on parameters θ and data x as
[Equation]
where T g refers to the group representation under which the data transforms (e.g., for translations, the pixel relabeling).
We are taking the convention that G is deﬁned by its action on θ , so we do not introduce an explicit representation on parameters.
The posterior is then equivariant under G if
[Equation]
Intuitively, this means that transforming the parameters with g is equivalent to transforming the observation with T g .
Here, J g is the Jacobian of the group transformation, which arises from the change of variables rule for probability densities (see App. A.1).
For translations | det J g | = 1 , but this is not true in general.
We propose to incorporate equivariances by taking the G -transformed parameters to ﬁducial values— i.e., to standardize the pose.
For our canonical example, this would mean translating the object to the center of the image.
In general, θ can also include parameters unchanged under such a transformation (e.g., the color of the object).
For each θ we denote the corresponding pose-standardized parameters as θ 0 .
These are related to θ by a group transformation g θ (the “pose” of θ ) such that g θ θ
= θ .
In this terminology, standardizing the pose means to take it to the group identity element e ∈ G .
3. P ROXY PARAMETERS FOR THE POSE
We are now presented with the basic problem that we resolve in this work.
If we had access to samples from p ( g θ | x ) , then a neural density estimator trained on pose-standardized data would sufﬁce for estimating the remaining parameters.
However estimating the pose itself would require an additional network trained on generic (non-pose-standardized) data.
The entire pipeline would therefore be limited by the latter network (as we demonstrate empirically in section 5.3).
To resolve this we make a compromise between precisely standardizing the pose and enabling parameter estimation: we transform the pose g θ to lie within a narrow band around e .
As we discuss in the following subsection, this gives rise to an iterative algorithm which takes successively better approximations to the parameters while transforming the pose to lie within the band.
We ﬁrst augment the parameter space to include an approximation (or proxy ) ˆ g ∈ G for the pose.
This is deﬁned using a kernel to blur the true pose, i.e., ˆ g = g θ (cid:15) for (cid:15) ∼ κ ( (cid:15) ) .
The kernel κ ( (cid:15) ) —a distribution over group elements (cid:15) ∈ G —speciﬁes the band to which we standardize the pose.
It should be chosen to be concentrated around e , and we furthermore choose it to be symmetric.
A natural choice is a simple Gaussian or uniform distribution.
Then ˆ g is a close approximation to g θ , and data x standardized with respect to ˆ g has pose within the band around e .
With this deﬁnition, p (ˆ g | θ )
= κ (cid:0) ( g θ ) − 1 ˆ g (cid:1) .
For translations, the group operation is addition and the identity element is the origin, so this discussion reduces to translating an image to place an object near the origin.
If p ( θ | x ) is equivariant under G , then it remains so when also conditioned on ˆ g (see App. A.2).
The latter distribution may therefore be expressed in terms of ˆ g -standardized quantities,
[Equation]
On the last line we deﬁned θ (cid:48) and x (cid:48) , and dropped the constant argument ˆ g − 1 ˆ g = e .
Hence, we can express the posterior conditioned on ˆ g in terms of the ˆ g -standardized posterior—which is much easier to estimate.
Importantly, if the equivariance is not exact we can still standardize the pose, but we must explicitly condition the ˆ g -standardized posterior on ˆ g to account for equivariance-breaking effects.
3. T RAINING AND INFERENCE
In the general case, we train a neural density estimator q ( θ (cid:48) | x (cid:48) , ˆ g ) to approximate the posterior with ˆ g -standardized data.
To do so, we augment the NPE maximum likelihood objective (1) with sampling over ˆ g during training,
[Equation]
(Note that the logarithm causes the Jacobian in (5) to drop out upon taking the gradient, since it is independent of q .)
We have explicitly conditioned q on ˆ g to allow for approximate equivariances.
Exact equivariances can be strictly enforced by dropping this.
Furthermore, for approximate equivariances, it may be convenient to omit the transformation ˆ g − 1 applied to θ , which is permitted with the ˆ g conditioning.
Given a trained q ( θ (cid:48) | x (cid:48) , ˆ g ) and a ﬁxed p (ˆ g | θ ) , we can self-consistently perform inference using Gibbs sampling.
At inference time, samples from the joint posterior p ( θ, ˆ g | x ) can be obtained by alternately sampling θ and ˆ g , conditionally on the other variable and x ,
[Equation]
This 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 κ ( (cid:15) )
the chain is both 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 sufﬁciently 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.
Convergence of the chain also informs our choice of κ ( (cid:15) ) .
For very wide κ , ˆ g is uninformative about the pose, so sampling (8) reduces to sampling from p ( θ | x ) , and GNPE reduces to NPE.
For very narrow κ , by contrast, the pose is highly constrained but many Gibbs iterations are needed to converge.
The width of κ is an additional hyperparameter, which should be chosen to accommodate this practical trade-off between convergence speed and accuracy.
As a rule of thumb, the standard deviation of a typical pose posterior is a suitable starting point.
In practice, we can obtain N samples in parallel by constructing an ensemble of N Markov chains according to equations (7) and (8).
We initialize these chains using samples from a second neural density estimator q init ( g θ | x ) , trained using standard NPE.
GNPE iterations then yield a sequence of sample sets { θ ( i ) j } Ni =1 , j = 0 , 1 , 2 , . .
, each of which represents a distribution Q j ( θ | x ) over parameters.
Assuming a perfectly trained network, a GNPE iteration applied to sample set j yields an updated distribution,
[Equation]
The “ ∗ ” symbol denotes a group convolution, the “ ¯ ∗ ” symbol a combination of marginalization and group convolution.
See App. A.3 for a derivation and further details.
In this notation, the initialization of the pose g θ in iteration 0 with q init simply means setting Q 0 ( ·| x ) ¯ ∗ κ = q init ( ·| x ) ∗ κ .
The true posterior p ( θ | x ) is clearly a ﬁxed point of this sequence, with the number of iterations to convergence determined by κ and the accuracy of the initialization network q init .
With κ just wide enough so that Q 0 ( ·| x ) ¯ ∗ κ ≈ p ( ·| x ) ¯ ∗ κ for a typical observation, a single GNPE iteration sufﬁces.
Alternatively, a narrower κ can be chosen to improve accuracy of q ( θ (cid:48) | x (cid:48) , ˆ g ) at the cost of additional iterations.
T OY EXAMPLE : DAMPED HARMONIC OSCILLATOR
As our ﬁrst example, we 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 undamped 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. B.1).
The model is constructed such that the posterior is equivariant under translations in τ ,
[Equation]
so we take τ to deﬁne the pose.
Note that the equivariance for this simple toy model is exact, but it could be easily turned into an approximate equivariance by, e.g., introducing a τ -dependent noise term.
The prior p ( τ ) extends from − 5 s to 0 s, so an NPE inference network would be tasked with interpreting data from oscillators excited throughout this range of times.
[Figure] For GNPE, we shift the data to align the pose within a narrow range around τ = 0 (Fig. 2a).
To do so we ﬁrst deﬁne a proxy ˆ τ = τ + (cid:15) for the pose, where (cid:15) ∼ κ ( (cid:15) ) is drawn from a zero-centered
Gaussian kernel with standard deviation 0.1 s. We then train a neural density estimator q ( θ (cid:48)
| x (cid:48) ) to approximate p ( θ (cid:48)
| x (cid:48) ) , where θ (cid:48) ≡ ( ω 0 , β, τ − ˆ τ ) and x (cid:48) ≡ T − ˆ τ x are pose-standardized.
We take q ( θ (cid:48) | x (cid:48) ) to be diagonal Gaussian, matching the known form of the posterior.
Note that by dropping the explicit conditioning of q on ˆ τ [cf.
(6)] we enforce an exact equivariance.
For each experiment, we train until the validation loss stops decreasing.
We also train a neural density estimator q init ( τ | x ) with standard NPE on the same dataset to generate initial GNPE seeds.
To generate N posterior samples we proceed as follows:
Sample τ
( i ) ∼ q init ( τ | x ) , i = 1 , . . .
, N ; 2.
Sample (cid:15) ( i ) ∼ κ ( (cid:15) ) , set ˆ τ ( i ) =
τ ( i ) + (cid:15) ( i ) , and time-translate the data, x (cid:48)
( i ) =
T − ˆ τ ( i ) x ; 3.
Sample θ (cid:48)
( i ) ∼ q ( θ (cid:48) | x (cid:48)
( 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.
B.2.
We evaluate GNPE on ﬁve simulations by comparing inferred samples to the ground-truth posteriors using the c2st score (Friedman, 2004; Lopez-Paz & Oquab, 2017).
This score corresponds to the test accuracy of a classiﬁer trained to discriminate samples from the target and inferred distributions, and ranges from 0 .
5 (best) to 1 .
(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.
B.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. 2b).
This underscores the importance of using equivariance-aware inference methods.
To visualize the effect of pose standardization, we show the weights of several ﬁlters from the ﬁrst layer of the fully trained GNPE and NPE embedding (Fig. 2c).
The NPE network is trained to interpret signals from oscillators excited at arbitrary τ , while GNPE can focus on signals starting around τ =
(up to a small κ perturbation).
This is also reﬂected in a reduced effective dimension of the input data to GNPE (Fig. 2d).
G RAVITATIONAL - WAVE PARAMETER INFERENCE
Gravitational waves—propagating ripples of space and time—were ﬁrst 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 scientiﬁc 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. C.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 h
I ( θ ) in each of the LIGO/Virgo interferometers
I = H , L , V .
The signals on Earth are very weak and embedded in detector noise n I .
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 = { h I ( θ ) + n I } 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 dimensionality 5 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 ﬁrst time performance indistinguishable from “ground truth” stochastic samplers—at drastically reduced inference times.
5. E QUIVARIANCES OF SKY POSITION AND COALESCENCE TIME
We consider the analysis of BBH systems.
Included among the parameters θ are the time of coalescence t c (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 t I of the signal in each of the interferometers.
6 Our priors (based on the precision of detection pipelines) constrain t I 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 t I —equivalently, ( t c , α, δ ) —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,
[Equation]
Group elements g abs ∈ G abs act by uniform translation of all t I , whereas g rel ∈ G rel act by individual translation of t
L and t V .
We work with data in frequency domain, where time translations act by multiplication, T g x I =
e − 2 πif ∆ t I x I .
Absolute time shifts correspond to a shift in t c , and are an exact equivariance of the posterior, p ( g abs θ | T g abs x )
= p ( θ | x ) , where | det J g abs | = 1 .
Relative time shifts correspond to a change in ( α, δ ) (as well as t c ).
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.
5. A PPLICATION OF GNPE
We use GNPE to standardize the pose within a band around t I = 0 .
We consider two modes deﬁned by different uniform blurring kernels.
The “accurate” mode uses a narrow kernel κ narrow =
[Figure]
U [ − 1 ms , 1 ms] n I , whereas the “fast” mode uses a wider kernel κ wide =
U [ − 3 ms , 3 ms] n I .
The latter is intended to converge in just one GNPE iteration, at the cost of having to interpret a wider range of data.
We deﬁne the blurred pose proxy
ˆ g
I ≡ t I + (cid:15) I , where (cid:15)
I ∼ κ ( (cid:15) I ) .
We then train a conditional density estimator q ( θ (cid:48) | x (cid:48) , ˆ g rel ) , where θ (cid:48)
=
ˆ g − 1 abs θ and x (cid:48)
=
T ˆ g − 1 x .
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 additionally train a density estimator q init ( { t I } I
=H , L , V | x ) using standard NPE to infer initial pose estimates.
We stress that the difﬁculty of the inference problem (high data dimensionality, signiﬁcant noise levels, complex forward model) combined with high accuracy requirements to be scientiﬁcally useful requires careful design decisions.
In particular, we initialize the ﬁrst 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 ﬂow (Durkan et al., 2019) to model the complicated GW posterior structure.
See App.
C.2 for details of network architecture and training.
5. R ESULTS
We evaluate performance on all eight BBH events from the ﬁrst Gravitational-Wave Transient
Catalog (Abbott et al., 2019) consistent with our prior bounds (component masses greater than 10 M (cid:12) ).
We generate reference posteriors with the standard LIGO/Virgo code LALI NFERENCE (Veitch et al., 2015), which uses MCMC.
We quantify the deviation between our inferred samples and the reference samples with the c2st score.
We compare performance against two baselines, standard NPE and a modiﬁed approach that partially standardizes the pose (‘chained NPE’).
For the latter, we use the chain rule to decompose the posterior,
[Equation]
(12) where λ =
( t c , α, δ ) are the parameters required to determine the pose, and φ ⊂ θ collects the remaining 12 parameters.
We use standard NPE to train a ﬂow q ( λ | x ) to estimate p ( λ | x ) , and a ﬂow q ( φ | x (cid:48) , λ ) to estimate p ( φ | x, λ ) .
The latter ﬂow is conditioned on λ , which we use to standardize the pose of x .
In contrast to GNPE, as noted in section 3.3, this approach is sensitive to the initial pose estimate q ( λ | x ) , which limits accuracy (Figs. 3 and C.7).
We note that all hyperparameters of the ﬂow and training protocol (see App.
C.2) were extensively optimized on NPE, and then transferred to GNPE without modiﬁcation, 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. 3).
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 ﬁnd
MCMCNPEGNPE m 2 [
M (cid:12) ]
[Figure]
1 [M (cid:12) ]
[-8 1 .2 .θ
-]
[Figure] [M (cid:12) ] [Figure]
θ 0 .
[Figure] MCMCNPEGNPE m
[Figure] M (cid:12) ] [-0 . 8 1 .2 .θ
-][Figure]
[-40 50 60 70 m 
-][Figure]
(cid:12) ] [Figure]
[-8 1 .2 .θ
 20 25 30 35 m 2 [M-] (cid:12)
[Figure]
that inaccuracies of that baseline can be almost entirely attributed to the initial pose estimate (Fig. C.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.
C.7). Both GNPE models signiﬁcantly 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 .
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.
The quantitative agreement between GNPE and MCMC (Fig. 3) is visible from the overlapping marginals for all parameters displayed.
NPE, by contrast, deviates signiﬁcantly 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.
C ONCLUSIONS
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 scientiﬁc forward models and any inference network architecture.
The only architectural modiﬁcation occurs for approximate equivariances, where additional context variables must be included to describe the pose proxy.
GNPE 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 efﬁciency and accuracy.
For the motivating problem of GW parameter estimation, GNPE achieves for the ﬁrst time rapid amortized inference with results virtually indistinguishable from MCMC (Anonymous et al., 2021).
This is an extremely challenging “real-world” scientiﬁc problem, with high-dimensional input data, complex signals, and signiﬁcant 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.
E THICS S TATEMENT
Our method is primarily targeted at scientiﬁc 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.
R EPRODUCIBILITY S TATEMENT
The experimental setup for the toy model is described in App.
B.2.
We also provide the code at https://tinyurl.com/wmbjajv8.
The setup for GW parameter inference is described in App.
C.2.
We will release a python package including scripts for the experiments carried out in section 5 upon acceptance and will also be happy to share them with the reviewers anonymously during the rebuttal period as suggested on the ICLR
FAQ
A D ERIVATIONS A. E QUIVARIANCE RELATIONS
Consider a system with an exact equivariance under a joint transformation of parameters θ and observations x ,
[Equation]
An invariant prior fulﬁlls the relation
[Equation]
where the Jacobian J g arises from the change of variables rule for probability distributions.
It is required for the normalization to remain intact,
[Equation]
Note that we neglect boundary effects here.
A similar relation holds for an equivariant likelihood,
[Equation]
Again
[Equation]
the Jacobian is required for the normalization,
[Equation]
[Equation]
An invariant prior and an equivariant likelihood further imply for the evidence p ( x )
[Equation]
[Equation]
Combining an invariant prior with an equivariant likelihood thus leads to the equivariance relation
[Equation]
for the posterior, where we used Bayes’ theorem and equations (15), (17), and (19).
A. E QUIVARIANCE 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 (5).
With the deﬁnition p
(ˆ g | θ )
= κ (cid:0) ( g θ ) − 1 ˆ g (cid:1) from section 3.3, the p (ˆ g | θ ) is equivariant under joint application of h ∈ G to ˆ g and θ ,
[Equation]
Here, we used that by deﬁnition a group element h directly alters the pose, g hθ = hg θ .
This implies, that p (ˆ g | x ) is equivariant under joint application of h and T h ,
[Equation]
(23) in the second step we used p (ˆ g | θ, x ) =
p (ˆ g | θ ) .
From these relations, the equivariance relation used in equation (5) follows,
[Equation]
A. I TERATIVE INFERENCE AND CONVERGENCE
GNPE leverages a neural density estimator of the form q ( θ | x (cid:48) , ˆ 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 (9), which states how a distribution Q j ( θ | 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,
[Equation]
For a given proxy sample ˆ g , a (perfectly trained) neural density estimator infers θ with
[Equation]
where we used p (ˆ g | θ )
= κ (( g θ ) − 1 ˆ g ) .
Combining (25) and (26), the updated distribution over θ samples reads
[Equation]
Here, κ ( − ) denotes the reﬂected kernel, κ ( − ) ( g )
= κ ( g − 1 ) ∀ g .
Since we choose a symmetric kernel in practice, we use κ = κ ( − ) in (9) [Equation]
B T OY E XAMPLE B. F ORWARD MODEL
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 inﬁnitely short pulse.
The time evolution of that system is governed by the differential equation
[Equation]
(28) 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
[Equation]
[Equation] This equation describes a deterministic, injective mapping between parameters θ = ( ω 0 , β, τ ) and a time series observation x ,
[Equation]
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
[Equation]
) instead.
We sample δθ from an uncorrelated Gaussian distribution
[Equation]
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 speciﬁed 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
[Equation]
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.
B. I MPLEMENTATION
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 ( θ (cid:48) | x (cid:48) ) and q init ( τ | x ) .
For further hyperparameters, we use the defaults of the sbi package (Tejero-Cantero et al., 2020).
B. R ESULTS
For all methods, we compute the average classiﬁer two-sample test score (c2st) based on 10, samples from the estimated and the ground truth posterior for ﬁve different simulations.
We then average the accuracy across 10 different seeds.
C G RAVITATIONAL WAVE PARAMETER INFERENCE C. F ORWARD MODEL AND AMORTIZATION
The forward model mapping binary black hole parameters θ (Tab. C.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) S n , drifts over the duration of an observing run.
To fully amortize the computational cost, we use a variety of different PSDs S n in training, and additionally condition the inference network on S n .
At inference time, this enables instant tuning of the inference network to the PSD estimated at the time of the event, see Anonymous et al.
(2021) for details.
Since this conditioning on
S n has no effect on the GNPE algorithm outlined in this work, we keep it implicit in all equations.
C. N ETWORK ARCHITECTURE AND TRAINING
The inference network consists of an embedding network, that reduces the high dimensional input data to a 128 dimensional feature vector, and the normalizing ﬂow, that takes this feature vector
[Table]
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 .
Hz
[Table]
and PSD information (10 46 · S n ) − 1 / 2 with the same binning.
This results to a total of (3 · 8,033) =
24,099 real input bins per detector.
The ﬁrst module of the embedding network consists of a linear layer per detector
[Table]
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 ﬁlter 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 ﬂow 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 ﬂow (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 hashidden layers and 1 . 31 · 10 8 (for two detectors) or 1 . 42 · 10 8 (for three detectors) learnable parameters.
We train the inference network with a data set of 5 · 10 6 waveforms with parameters θ sampled from the priors speciﬁed in table C.1, and reserve 2% of the data for validation.
We pretrain the network with learning rate of 3 · 10 − 4 for 300 epochs with ﬁxed PSD, and ﬁnetune 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.
C. R ESULTS
The c2st scores between inferred posterior and the MCMC reference shown in Fig.
3 are computed using the code and default hyperparameters of Lueckmann et al.
For each event, we compute the c2st score of 10,000 samples for inferred and target posterior.
Fig.
3 displays the mean of the score across 5 different sample realizations, Fig.
C.1 additionally shows the corresponding standard deviation.
For technical reasons we use only 12 of the 15 inferred parameters; speciﬁcally we omit the geocentric time of coalescence t c (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).
[Figure]
GW
GW
GW
GW
GW
GW
GW
GW
[Figure]
GNPE with a narrow kernel consistently outperforms the baselines, which is in accordance with Fig.
[Figure]
κ
[Figure]
[Figure]
κ
[Figure]
GW GW
[Figure]
GW
GW
GW
GW GW [Figure] MCMC
GNPE    m[ M (cid:12) ] [Figure] [Figure] [Figure]
[Figure] θ [Figure] 1 .2 .θ
{+t H [ m s ] 32 40 48 56m 
+}
[M (cid:12) ] ˆ t H [ m s ]
[Equation]
[Equation]
MCMCchained NPE, NPE pose chained NPE, oracle pose
[Figure] m [Figure] [ M
[Figure]
[Figure]
[Figure]
t H [ m s ]
[Figure]
M [Figure] ] t L [ m s ] [Figure]
32 .35 .37 .40 .t
H [ms]     t L [ m s ] NPE pose oracle pose      ˆ t H [ms]    ˆ t L [ m s ] NPE pose proxy oracle pose proxy