Group equivariant neural posterior estimation
Authors: Maximilian Dax
Stephen R Green
Jonathan Gair
Michael Deistler
Bernhard Schölkopf
Jakob H. Macke
Venue: ICLR 2022 Poster (2022)
Keywords: simulation-based inference
likelihood-free inference
machine learning for science
equivariances
group transformations

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 gravitational-wave observations. We show that GNPE achieves state-of-the-art accuracy while reducing inference times by three orders of magnitude.

Published as a conference paper at ICLR 2022
GROUP EQUIVARIANT
NEURAL POSTERIOR ESTIMATION
Maximilian Dax
Stephen R. Green
Max Planck Institute for Intelligent Systems Tübingen, Germany
Max Planck Institute for Gravitational Physics Potsdam, Germany
maximil ian, www.tuebbigen..nno...m www.eeiiinnnn.g.........mm de
st stephen. Eeephenggreenaeei www.ggeeennni.enne wat.aeee.pp.... mpg.de
Jonathan Gair
Michael Deistler
Max Planck Institute for Gravitational Physics Potsdam, Germany
Machine Learning in Science, University of Tübingen Tübingen, Germany
Bernhard Schölkopf
Jakob H. Macke
Max Planck Institute for Intelligent Systems Tübingen, Germany
Max Planck Institute for Intelligent Systems
Machine Learning in Science University of Tübingen Tübingen, Germany
ABSTRACT
Simulation- based inference with conditional neural density estimators a powerful approach solving inverse problems in science. However, these methods typically treat the underlying forward model as 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-calles group equivariant neural posterior estimation (GNPE)-is based on self-consistently standardizing the "pose" of the data while estimating the posterior over parameters. is architecture-indeeedder and applies both to exact and approximate equivariances. As real-world application, amortized inference of astrophysical binary black hole systems from gravitational- wave observations. We show that GNPE achieves taae-of-the-an accuracy while reducing inference times by three orders of magnitude.
GNPE for
INTRODUCTION
Bayesian inference provides a means of characterizing a system by comparing models against data. Given forward model or likelihood p(x]0) for data x described by parameters 0, and prior p(0), the Bayesian posterior is proportional to the product, p(0[z) o p(x)0)p((). Sampling techniques such as Markov Chain Monte Carlo (MCMC) can be used to build up posterior distribution provided the likelihood and prior can be evaluated.
For models with intractable or expensive likelihoods (as often arise in scientific applications) simulation-baze (or likelihood-free inference methods offer powerful alternative (Cranmer et al.) 2020). In particular, neural posterior estimation (NPE) (Papamakarios Murray] 2016) uses ex- pressive conditional density estimators such normalizing flows (Rezende & Mohamed] 2015) Papamakarios et al.) 2021) to build surrogates for the posterior. These are trained using model simulations x ~ p( 0). and allow for rapid sampling for any p(x), thereby amortizing training costs across future observations. NPE and other density-estin estimation methods for simulation. based inference (Gutmann & Corander 2016 Papamakarios et al.) 2019 Hermans al.) (0200 have been reported to be more simulation- methods such Approximate Bayesian Computation (Sisson et al. 2018
(Lueckmann et al. 202 1) than classical likelihood-fii

---

Published as a conference paper at ICLR 2022
wmm
By standardizing the source sky position, GW signal can be made to arrive at the same time in all three LIGO/Virgo detectors. However, since this also changes the projection of the signal onto the detectors, it defines only an approximate equivariance. Nevertheless, our proposed GNPE algorithm simplifies inference by simultancously inferring and standardizing the incident direction.
Figure 1:
Virgo
Hanford
www.
standardized www.ae
Livingston
Training inference network for any p(x) can nevertheless present challenges due the large number of training samples and powerful networks required. The present study is motivated by the problem of gravitational-wan wave (GW) data analysis. Here the task infer properties of astrophysical black-hole mergers based 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- -asee methods (Veitch et al.) Ashton Ashton et al. 2019 The GW posterior, however, is equivarian Jnder 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. TT If we could constrain these parameters 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 date annd simplify analysis. The obvious roadblock is that the pose is contained in the set of parameters 0 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 proxy for the pose- blurred version- on which one conditions the posterior. The pose of the data is then transformed based on the proxy, placing it in band about the standard value, and resulting in an easier inference task. Finally, the joint posterior over 0 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 network architectures. For GWs, for example, we use specialized embedding networks to extract signal waveforms from frequency- domain data, well as expressive normalizing flows to estimate the neiterr of which is ttraighfoowwe 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-iddependenn it applies whether equivariances are exact or approximate, and allows for arbitrary equivariance relations between parameters and data.
We discuss related work in Sec. fec.2las and describe the GNPE algorithm in Sec. 3 In Sec. 4 4we apply GNPE 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 al.) 2015). The present paper describes the GNPE method which developed for GW analysis (Dax et 202 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. (221I)
RELATED WORK
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
In physics, the term "covariant" is frequently used instead of "cquivariant".
We adopt the language from computer vision by Jaderberg al. (2015).
2

---

Published as a conference paper at ICLR 2022
conflict with design considerations such 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 transfor- mations 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 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. (020) by considering "gray-box" setting. GNPE allows us to incorporate hige-level domain knowled ledge 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. (0018) is by constructing (Fearnhead & Prangle, 2012) or learning (Jiang et al. 2017 Chen et al. 2021) equivariant summary statistics (æ) which are used as input to the inference algorithm instead of the raw data æ. How- ever, 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- lcamming techniques for amortized GW parameter inference (Gabbard et al. 2019 Chua & Vallisneri 2020, Green & Gair 2021 Delaunoy 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 [Dx et al. (2021) the only one treat the full amortized parameter inference problem with accuracy matching standard methods.
METHODS
3.1 NEURAL POSTERIOR ESTIMATION
NPE (Papamakarios & Murray 2016 Greenberg et al. (019) is simulation- based inference method
that directly targets the posterior. Given a dataset of prior parameter samples (0)) ~ p(0) and
corresponding model simulations ((()) (z2(0((). it trains neural density estimator g(0 x) to
estimate ()) This is achieved by minimizing the loss
CNPE Epco) pp( [- log ( x=]
across the dataset of (()) pairs. This maximum likelihood objective leads to recovery of p(0
q(f x)) is sufficiently flexible. Normalizing flows (Rezende & Mohamed 2015 Durkan et al. 2019)
a particularly expressive class of conditional density estimators commonly used for NPE.
NPE amortizes inference: once q(([z) x)) is trained, inference is very fast for any observed data 2o. 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 exploit improve learning.
3.2 EQUIVARIANCES UNDER TRANSFORMATION GROUPS
In this work we describe a generic method to incorporate equivariances under joint transformations of 0 and x into NPE. A typical example arises when inferring the position of object from image data. In this case, if we spatially translate an image by some offset ď effected by relabeling the pixels- then the inferred position d should also transform by by addition to the position coordinates A. Translations are composable and invertible, and there exists trivial identity translation, 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 keep the discussion general.
3

---

Published as a conference paper at ICLR 2022
For transformation group G, we denote the action of g G parameters and data as
0 - g+g
(2)
I Tzz.
(3)
Here, Tg refers to the group representation under which the data transform (c.g., for image translations, the pixel relabeling). We adopt the natural convention that G is defined by its action 0, do not introduce an explicit representation on parameters. The posterior distribution p(0|z;) is said to be equivariant under G if, when the parameter and data spaces are jointly G-transformed, the posterior is unchanged, i.c.,
p(0|=) T(g))TT)) Jyl
Vg E G.
The right-hand side comes from the change-of-variabe 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. A.I
rule. For translations the Jacobian J, has
Our goal is to use equivariances to simplify the data to G-transform such that 0 is taken to a fiducial value. For the image example, this could mean translating the object of interest to the center. In general, 0 can also include parameters unchanged under G' (e.g., the color of the object), so we denote the corresponding standardized parameters by 0g. These are related to by a group transformation denoted gº, such that 000n 0. We refer to go as the "pose" of and standardizing the pose means to take it to the group identity element E G. Applying - to the data space effectively reduces its dimensionality, making it easier interpret for 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 nor 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 non-invariant prior, or if the forward model is itself non-equivariant.
3.3 GROUP EQUIVARIANT NEURAL POSTERIOR ESTIMATION
We are now presented with the basic problem that we resolve in this work: how to simultancously infer the pose of signal and use that inferred pose to standardize (or align) the data so to simplify the analysis. This is a circular problem because one cannot standardize the pose (contained in model parameters 0) 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 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 9 G. These "pose proxies" are defined using kernel to blur the true pose, i.e., 9 = g for k((); then p(ğ)0) = k (()")-10). The kernel x(e) 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 s(e) include Gaussian and uniform distributions. For translations, the pose proxy is simply the true position with additive noise.
Consider now the posterior distribution p(0, glæz) over the expanded parameter space. Our iterative algorithm comes from Gibbs sampling this distribution (Roberts & Smith 1994) (Fig. 2 i.c., alternately sampling 0 and ğ. conditional on the other parameter and x,
0 ~ (f 9).
(5)
9 ~ p(0l=,)) 0).
(6)
The second step just amounts to blurring the pose, since p(ğ) (r, = r(pl0) x (1g)) The key first step uses neural density estimator that is trained taking advantage of a standardized pose.
For an equivariant posterior, the distribution (5) can be rewritten as (App. A.2)
p(0|. 9) 10T g 9) det Jn = (f(((x)) det JT.
For the last equality we defined 0' --- and x' T 1r, and we dropped the constant argument و-19 This expresses p00 pr, .) in terms of the j-standardized data which is much easier to
4

---

Published as a conference paper at ICLR 2022
Figure 2: We infer p(g" [[/z) with Gibbs sampling by alternately sampling (1) g p(g" lz. 9) (blue) and (2) 9 ~ p(g]a, 10. (orange). For (1) we use a density estimator q(g (z)) for (2) the def- inition 9 = 9 - E e k(e). Pose standardization with T3 is only allowed due to conditioning on g. Increasing the width of y accelerates conver- gence (due to larger steps in parameter space), at the cost of 9 being a worse approximation for g", and therefore pose alignment being less effective.
= -'.02 0.2]
= 1111,11
plg" -g(=)
g
g"
estimate. We train neural density estimator q('")x') ) to approximate this, by minimizing the loss, LONPE Bp(o) (pt [- log q (ø ([TT-
(8)
With a trained (0),,,
~ p(0 1 g)
0 = g0', 0' ~ q(0 IT).
The estimated posterior is equivariant by construction (App. A.3
For an approximately-equvvriinn nevertheless able to the conditioning on I to approximately align æ. We therefore train neural density estimator q(0 x g). by minimizing the loss
posterior, ( cannot be transformed to be independent of g. We
CUNE = (p) Ep 0) p((10) [- [olog q (0(Tg-12,))]
(10)
In general, one may have combination of exact and approximate equivariances (see, e.g., Sec. 5
3.4 GTBBS CONVERGENCE
The Gibbs-sampling procedure constructs Markov chain with equilibrium distribution p(0, g(rr). For convergence, the chain must be transient, aperiodic and irreducible (Roberts & Smith 1994] Gelman ct al. 2013). For sensible choices of x<e) the chain is transient and aperiodic by construction. Further, irredueibility 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 required to ensure (approximately) independent samples. By marginalizing over (i.e., ignoring it) we obtain samples from the posterior p(0)zr). as desired
Convergence of the chain also informs our choice of k(e). For wide rá(e), only few Gibbs iterations needed to traverse the joint posterior p(0.g)z), z)), whereas for narrow k<e) many steps are required (Fig. 2) In the limiting case of n(e) a delta distribution (i.c., no blurring) the chain does not deviate from its initial position and therefore fails to converge Conversely, narrower k:(c) better constrains the pose, which improves the accuracy of the density estimator. The width of r should be chosen based on this practical trade-off between speed and accuracy; the standard deviation of typical pose posterior is usually good starting point.
In practice, obtain N samples in parallel by constructing ensemble of N Markov chains. We initialize these using samples from a second neural density estimator qunit (g [c), trained using standard NPE. Gibbs sampling yields a sequence of sample sets (0f)) J 0, 1, 2, each of which represents distribution Q .r) over parameters. Assuming perfectly trained network, one iteration applied to sample set j yields an updated distribution.
[2((lz) &
@)+1(0/x) = p(([z)
p( * (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(0 (r) is clearly a fixed point of this sequence, with the number of iterations to convergence determined by and the accuracy of the initialization network Gimit-
In practice, this results only in approximate samples due to the asymptotic behaviour Gibbs sampling, and a potential mismatch between the trained and the targeted true posterior. *This also explains why introducing the pose proxy is needed all: GNPE would not work without it!
5

---

Published as a conference paper at ICLR 2022
Simulations
Performance
NN weights
Singular values
0.1
NPE
1.0
NPE NPE-CNN
0.08
NPE
NPE
0.9
0.00
10"
GNPE
GNPE
10-2
0.8
-0.08
10"
0.1
GNPF
0.7
GNPR
10"
0.6
0.00
10-5
0.5
-0.15
10-6
500
10000
20000
30000
1000
2000
100
200
300
Time (8)
Simolations
Dimectsion
Index singular value
Figure 3: Comparison of standard NPE (blue) and GNPE (orange) for the damped harmonic oscillator. a) Three sample imputs to the neural density estimators showing GNPE data are pose- pose-standardized. b) setiiaiiance-agnostte NPE, and is on par with NPE with convolutional embedding network (purple). c) Example filters from the first linear layer of fully trained networks. GNPE filters more clearly capture oscillatory modes. d) Singular values of the training data. Inputs to GNPE x' have smaller effective dimension than raw inputs x to NPE.
c2st score performance (best: 0.5, worst: 1.0): GNPE significantly outperforms
4 TOY EXAMPLE: DAMPED HARMONIC OSCILLATOR
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 wo, damping ratio 3, and time of excitation т. The time series x is therefore a damped sinusoid starting at (and zero before). Noise introduced via normally-distributer perturbation of the parameters 0 = (Wo, T), resulting in a Gaussian posterior p(0]z) (further details in App. C.I). The model is constructed such that the posterior equivariant under translations T.
(12) so we take T to be the pose. The equivariance of this model is exact, but it could easily be made approximate by, e.g., introducing r-dependent noise. The prior p(r) extends from -5 s to 0 for NPE, the density estimator must learn to interpret data from oscillators excited throughout this range.
p(wo, 3, r Ar) = p(wo, B, r|x),
For GNPE, we shift the data to align the pose near 0 using Gaussian kernel =o0 (0.1 s) (Fig. 3d). We then train a neural density estimator q('' (x) to approximate p(0 [00| where 6' = pose-standardized. We take ((''(z') 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 ginit ginir(T (z)) with standard NPE on the same dataset to generate initial GNPE seeds. To generate N posterior samples we proceed as follows:
(wa- -c) and x' (r++))
1. Sample 7() ~ qmnn (T/x) i I,...,N;
2. Sample (()) ~ k(e), set (0) = = (i)) + 00, and time-translate the data, x((i) = T 20));;;
3. Sample ('(() ~ q(0" ()),,, and undo the time translation (()) to obtain ((()
We repeat steps 2 and 3 until the distribution over 7 converges. For this toy example and our choice of x only one iteration is required. For further details of the implementation see App. C22
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 classifier trained to discriminate samples from the target and inferred distributions, and ranges from 0.5 (best) to 1.0 (worst). As basclines we evaluate standard NPE (i) with 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-con -qquiviri-iuiviiiinn embedding network), perform similarly well and far outperform standard NPE (Fig. 3p.. 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 small perturbation) is also reflected in in simplified filters in the first layer of the network (Fig, (3b) and a reduced effective dimension of the input data to GNPE (Fig. (d))

---

Published as a conference paper at ICLR 2022
5 DRAVITATIONAL-WAY PARAMETER INFERENCE
Gravitational waves propagating ripples of space and time were first detected in 2015, from the inspiral, merger, and ringdown of pair of black holes (Abbott et al. (016). Since that time, the two LIGO detectors (Hanford and Livingston) (Aasi et al., 2015) as well as the Virgo detector (Acernese) jet al., 2015) have observed signals from over 50 coalescences of compact binaries involving either black holes or neutron stars (Abbott et al.] 2019 2 d). Key scientific results from these observa- tions 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 mcasurement 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 8, including the component masses and spins, as well the space-time position and orientation of the system (Tab. D.I. Given these parameters, Einstein's '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 h0) in each of the LIGO/Virgo interferometers I= L, V. The signals on Earth are very weak and embedded in detector noise ng. In part have a tractable likelihood, the noise is approximated as additive and stationary Gaussian. The signal and noise models give rise to a likelihood p(z(0) for observed data (h)(0) + ny n1)I-HLL.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.
Ashton Ashton et al., (019). 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) 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 challenge (Gabbard al. 2019 Chua & Vallisneri 2020: Green & Gair 202T: Delaunoy et al. 2020) to obtain results of comparable accuracy and completeness classical samplers. We now show how GNPE can be used 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.
5.1 EQUIVARIANCES OF SKY POSITION AND COALESCENCE TIME
We consider the analysis of BBH systems. Included among the parameters 0 are the time of coalescence fe (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 ły of the signal in each of the interferometers Our priors (based on the precision of detection pipelines) constrain ł to a range of x 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 e-uivallienn,, (tc, a, 8) 8)-can can alternatively be interpreted as the pose of the data, and standardized using GNPE. The group G transforming the pose factorizes into direct product of absolute and relative time shifts,
G G= Gats X Grel-
(13)
Group elements Jabs E Gals act by uniform translation of all f,, whereas Ieel E Gzel act by individual translation of 4 and by. We work with data in frequency domain, where time translations act by multiplication, Tzæ = e 22iifftt xT Absolute time shifts correspond to shift in le, and are an exact equivariance of the posterior, p( g.. = p(0]z). Relative time shifts correspond to change in (a, 5) (as well as tc). This is only an approximate equivariance, since 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.
'In 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. We consider observations made in either n, 2 or interferometers.
7

---

Published as a conference paper at ICLR 2022
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 out- performs both NPE baselines, while being only marginally slower (1 it- eration ~ 2 s). With narrow ker- nel and 30 iterations (~ 60 s), we achieve c2st 0.55 across all events. Ø indicates the average across all eight events. For an alternative met- ric (MSE) see Fig. D.2
GW150014
GW170500
0.70
GW153012
CW170814
GWI20
GW120729
GW170823
0.65
62st
0.627
0.60
07 0.558
0.55
2 0.526
Ø 0.519
0.50
NPE
NPE (chained)
GNPE (wide x)
GNPE (narrow K)
5.2
APPLICATION OF GNPE
We use GNPE to standardize the pose within band around 4 0. We consider two modes defined by different uniform blurring kernels. The "accurate" mode uses a narrow kernel Knarrow = U[-1 ms, 1 ms]"/, whereas the "fast" mode uses wide kernel Kówide U[-3 ms, 3 ms]" 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 Ô. = t, + c1, where ci ~ r((r). We then train a conditional density estimator g081|z', Iee),, where 6' = 9ahe 9 and x' = -g-II. That is, we condition g on the relative time shift (since this is approximate equivariance) and 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 density estimator qiiic (f,,)(((H,H,.vll 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 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 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.
5.3 RESULTS
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 Mo). We generate reference posteriors with the LIGO/Virgo MCMC code LALINFERENCE (Veitch et al. 2013). We quantify the deviation between NPE samples and the reference samples using c2st.
We compare performance against two baselines, standard NPE and modified approach that partially standardizes the pose ("chained NPE"). For the latter, we use the chain rule to decompose the posterior,
x00xx) p(ó, A/x) = p( A) p(A[=),
(14)
where A = (lc, ar, 5) are the pose parameters and C collects the remaining 12 parameters.
We use standard NPE to train flow q( (A/ar) to estimate p(A[z), and flow q(0)x², A) to estimate p(ólz, z. A). The latter flow is conditioned on X, which we use to standardize the pose of In contrast to GNPF, this approach is sensitive to the initial pose estimate (A((z)), which limits accuracy (Figs. 4 and D.7). We note that all hyperparameters of the flow and training protocol (see App. ( 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-movt 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 GW1708 1708T4, 14, for which the initial pose estimate is inaccurate. Indeed, we find
8

---

Published as a conference paper at ICLR 2022
MCMC NPE GNPE
MOMC NPE GNPE
4 t % $ %
[Ma]]
1G 24
A
03
82
0.
4 +
[Ma]
[Ma]
81
[M
ma [Ma]
02
Figure 5: Corner plots for the GW events GW170809 (left) and GW170814 (right), plotting ID marginals on the diagonal and 90% credible regions for the 2D correlations. We display the two black hole masses m1 and m2 and two spin parameters 01 and 02 (note that the full posterior is 15-dimensional), NPE does not accurately reproduce the MCMC posterior, while accurate-mode GNPE matches the MCMC results well. For plot with all baselines Fig. D.4
that inaccuracies 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 this effect due 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. a recent benchmark by Lueckmann al. (021) 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-iee, 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, GW1 108009 and GW170814, in Fig. The quantitative agreement between GNPE and MCMC (Fig. 4 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 GW1 70823, deviations of NPE from the reference posterior less clearly visible.
CONCLUSIONS
We described GNPE, approach incorporate exact-and even approxinatee-eqqivaraaace under joint transformations of data and parameters into simulation- based inference. GNPE can be applied 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 ((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. (021). This is an extremely challenging "real-world" scientific problem, with bigh-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.

---

Published as a conference paper at ICLR 2022
ETHICS STATEMENT
Our method is primarily targeted 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 detections increases with more sensitive detectors in the future.
REPRODUCIBILITY STATEMENT
The experimental setup for the toy model is described App. We also provide the code at https:/tinyurl.cc cow. wwmjajv88 The setup for GW parameter inference is described in App. D.2 We
will release python package including scripts for the experiments carried out in section
ACKNOWLEDGMENTS
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.D. B This research has made use of data, software and/or web tools obtained from the Gravitational Wave Open Science Center http:://www..wwopenscctieee
a service of LIGO Laboratory, the LIGO Scientific Collaboration and the Virgo Collaboration. LIGO Laboratory and Advanced LIGO funded by the United States National Science Foundation (NSF) as well as the Science and Tech- nology Facilities Council (STFC) of the United Kingdom, the Max-Planck-Sociee (MPS), and the State of NiedersachsennGGrrmann 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 funded, through the European Gravitational Observa- tory (EGO). by the French Centre National de Recherche Scientifique (CNRS), the Italian Istituto Nazionale Fisica Nucleare (INFN) and the Dutch Nikhef, with contributions by institutions from Belgium, Germany, Greece, Hungary, Ireland, Japan, Monaco, Poland, Portugal, Spain. This material 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 al., 2019)) nflows (Durkan et al. 2020) and sbi (Tejero-Canter et al., 2020) for the implementation of our neural networks. The plots are generated with matpl ot ib (Hunter 2007) and Cha i CConsumer (Hinton 2016).
REFERENCES
J. Aasi et al. Advanced LIGO. Class. Quant. Grav., 32:074001, 2015. 9381/32/7/074001
doi 10.0080044/
B. P. Abbott et al. Observation of Gravitational Waves from Binary Black Hole Merger. Phys. Rev. Leut, 116(6):06110) 102, 2016. 10110.1100//hyseeel..r 16.061100
B. P. Abbott et al. A gravitational wave standard siren measurement of the Hubble constant. Nature, 551(7778):85-88) 88, 2017. ( [0.1003//ature24447
B. P. Abbott et al. GW 17017817 Measurements of neutron star radii and equation of state. Phys. Rev. Lett., 121(16):16110)) 101, 2018.
B. P. Abbott al. GWTCC1: A Gravitational- Wave Transient Catalog of Compact Binary Mergers Observed by LIGO and Virgo during the First and Second Observing Runs. Phys. Rev. X, 9()): 031040, 2019. 10,110.11103/hhyyeeex.0.0.333
R. Abbott et al. GWTC-2: Compact Binary Coalescences Observed by LIGO and Virgo During the First Half of the Third Observing Run. Phys. Rev. X, 11:021053, 2021a. doi do. 1tw.1000///yssx..1...22222
R. Abbott et al. Population Properties of Compact Objects from the Second LIGO- LIGO-Virgo Gravitational- Wave Transient Catalog. Astrophys. Lett., 913(1):L.7, 2021b. 10..84700772
10

---

Alvin J. K. Chua and Michele Vallisneri. ral networks for gravitational wave inference. 1ot Published as a conference paper at ICLR 2022
R. Abbott et al. Tests of general relativity with binary black holes from the second LIGO- LIGO-Virgo gravitational-wave transient catalog. Phys. Rev. D, 103(12):12002,, 2021c.
R. Abbott et al. GWTC-2.1: Deep Extended Catalog of Compact Binary Coalescences Observed by LIGO and Virgo During the First Half of the Third Observing Run. 8 2021d.
Acernese et al. Advanced Virgo: a second -enerraion interferometric gravitational wave detector. Class. Quant. Grav., 32(2):024001, 2015. doi
Gregory Ashton et al. BILBY: A user friendly Bayesian inference library for gravitational
astronomy. Astrophys. Suppl., 241(2): 2019. doi 100.8447/133884355/abbbb6
Atilim Güneş Baydin, Lei Shao, Wahid Bhimji, Lukas Heinrich, Lawrence Meadows, Jialin Liu, Andreas Munk, Saeid Naderiparizi, Bradley Gram-Haasen, Gilles Louppe, al. Etalumis: Bringing probabilistic programming to scientific simulators at scale. In Proceedings of the international conference for high performance computing, networking, storage and analysis, pp. 1-24, 2019.
Alejandro Bohé, Mark Hannam, Sascha Husa, Frank Ohme, Michael Pürrer, and Patricia Schmidt. PhenomPv2 technical notes for the LAL implementation. LIGO Technical Document, LIGO- TI 5006002-444, 2016. URL https: /tcc//cccc/ooooo 1igo. org/ o
Denis Boyda, Gurtej Kanwar, Sébastien Racanière, Danilo Jimenez Rezende, Michael Albergo, Kyle Cranmer, Daniel C Hackett, and Phiala E Shanahan. Sampling using su (n) gauge equivariant flows. Physical Review D, 103(7):074504, 2021.
Johann Brehmer, Gilles Louppe, Juan Pavez, and Kyle Cranmer. Mining gold from implicit models to improve likelihood free inference. Proceedings of the National Academy of Sciences, 117(10): 2020.
Yanzhi Chen, Dinghuai Zhang, Michael Gutmann, Aaron Courville, and Zhanxing Zhu. Neural approximate sufficient statistics for implicit models. In Ninth International Conference on Learning Representations 2021, 2021.
Learning Bayesian posteriors with neu- Phys. Rev. Lett., 124(4):04112,, 2020.
Taco Cohen and Max Welling. Group equivariant convolutional networks. In International conference machine learning, pp. 2990-2999, PMLR, 2016
confrree
Kyle Cranmer, Johann Brehmer, and Gilles Louppe. The frontier of simulation- inference. Proceedings of the National Academy of Sciences, 117(48):0055-33002 2020.
Maximilian Dax, Stephen R. Green, Jonathan Gair, Jakob H. Macke, Alessandra Buonanno, and Bernhard Schölkopf. Real-Time Gravitational Wave Science with Neural Posterior Estimation. Phys. Rew Leut, 127(24)2241033, 2021. doi T0/hhoss//yysseeeeet.... 127.2444110
Arnaud Delaunoy, Antoine Wehenkel, Tanja Hinderer, Samaya Nissanke, Christoph Weniger, An- drew R Williamson, and Gilles Louppe. Lightning-fast gravitational wave parameter inference through neural amortization. Third Workshop Machine Learning and the Physical Sciences (NeurIPS 2020), 2020.
Conor Durkan, Artur Bekasov, lain Murray, and George Papamakarios. Neural spline flows. Advances in Neural Information Processing Systems, 32:751 12:7511-75522 11-522 2019.
Conor Durkan, Artur Bekasov, Iain Murray, and George Papamakarios. nflows: normalizing flows in PyTorch, November 2020. URL httpp:///ol..rr//11...
42962877
Benjamin Farr, Evan Ochsner, Will M. Farr, and Richard O' Shaughhesssy. A more effective coordinate system for parameter estimation of precessing compact binaries from gravitational waves. Phys. Rev. 1), 90(2):024018, 2014. doi 1tt...111033hyykeee.0000000000

---

Published as a conference paper at ICLR 2022
Paul Fearnhead and Dennis Prangle. Constructing summary statisties for approximate bayesian computation: remi-automatic approximate bayesian computation. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 74(3):419-474, 2012.
Jerome Friedman. On multivariate goodness-of-fit and two-sample testing. In Conference Statistical Problems in Particle Physics, Astrophysics and Cosmology, 2004.
Hunter Gabhard, Chris Messenger, Ik Siong Heng, Francesco Tonolini, and Roderick Murray-Smith. Bayesian parameter estimation using conditional variational autoencoders for gravitational-wave astronomy, 2019.
Andrew Gelman, John B Carlin, Hal S Stern, David Dunson, Aki Vehtari, and Donald B Rubin. Bayesian data analysis. CRC press, 2013.
Stephen R. Green and Jonathan Gair. Complete parameter inference for GW LSW150914 using deep learning. Mach. Learn Sci. Tech., 2(3)03LTT1, 2021. doi [0111111666232/1//aaaaellll
David Greenberg, Marcel Nonnenmacher, and Jakob Macke. Automatic posterior transformation for likelihood-free inference. In International Conference on Machine Learning, pp. 2404-2414. PMLR, 2019.
Michael U Gutmann and Jukka Corander. Bayesian optimization for likelihood-free inference of simulator-based statistical models. The Journal of Machine Learning Research, 17(1):45564302, 2016.
Mark Hannam, Patricia Schmidt, Alejandro Bohé, Leila Haegel, Sascha Husa, Frank Ohme, Geraint Pratten, and Michael Pürrer. Simple Model of Complete Precessing Black-Hole-Binary Gravita- tional Waveforms. Phys. Rev. Lett., 111110, 2014. doi TO. 103/Physeeeleeett - 113.151101
Joeri Hermans, Volodimir Begy, and Gilles Louppe. Likelihood-free meme with approximate likelihood ratios. In Proceedings of the 37th International Conference on Machine Learning. volume 98 of Proceedings of Machine Learning Research. PMLR, 2020.
S. R. Hinton. ChainConsumer. The Journal of Open Source Software, 1:00045, August 2016. 100/
D. Hunter. Matplotlib: A 2d graphics environment. Computing in Science Engineering, 9()):
90-95, 2007. TO. 109/MCSE.2007.55
Max Jaderberg, Karen Simonyan, Andrew Zisserman, et al. Spatial transformer networks. Advances in neural information processing systems, 28:2017-200 2025, 2015.
Bai Jiang, Tung-yu Wu, Charles Zheng, and Wing H Wong. Learning summary statistic for ap- proximate bayesian computation via deep neural network. Statistica Sinica, pp. 1595-1618, 2017.
Sebastian Khan, Sascha Husa, Mark Hannam, Frank Ohme, Michael Pürrer, Xisco Jimúez Forteza, and Alejandro Bohé. Frequency-dom gravitational waves from nonprecessing black-hole binaries. II. A phenomenological model for the advanced detector era. Phys. Rev., D43(4)0440007, 2016. TO. 103/PhyyeeeD.....40000
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep con- volutional neural networks. Advances in neural information processing systems, 25:1097-1105. 2012.
David Lopez-Paz and Maxime Oquab. Revisiting classifier two-sample tests. In 5th International Conference on Learning Representations, ICLR, 2017.
Jan-Matthis Lueckmann, Jan Boelts, David Greenberg, Pedro Goncalves, and Jakob Macke. Bench- marking simulation-hased inference. In International Conference on Artificial Intelligence and Statistics, pp. 343-351. PMLR, 2021.
George Papamakarios and lain Murray. Fast E-free inference of simulation models with bayesian conditional density estimation. In Advances neural information processing systems, pp. 1028- 1036, 2016.
12

---

Published as a conference paper at ICLR 2022
George Papamakarios, David Sterratt, and lain Murray. Sequential neural likelihood: Fast likelihood- free inference with autoregressive flows. In Proceedings of the 22nd International Conference on Artificial Intelligence and Statistics (AISTATS), volume 89 of Proceedings of Machine Learning Research, pp. 837-848. PMLR, 2019.
George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. Journal of Machine Learning Research, 22(57):1-64 64. 2021. URL http: http:://mmm..ogg///aaaaaaa
19-1 028 html
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang. Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' Alché Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems 32, pp. 8024-8035, Curran Associates, Inc., 2019. URL /ttppr//ppppr.....m wtt.ppprssseerssccccceaaaaeeeeee 9015 pytorch an imperative style high performance deeee-leanttee pdf
Library.
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International conference on machine learning, pp. 1530-1538. PMLR, 2015.
Gareth O Roberts and Adrian FM Smith. Simple conditions for the convergence of the gibbs sampler and metropolis- hhastiiss algorithms. Stochastic processes and their applications, 49(2):207-216 1994.
Scott A Sisson, Yanan Fan, and Mark Beaumont. Handbook of approximate Bayesian computation. CRC Press, 2018.
Alvaro Pgjero-Cantero. Jan Boelts, Michael Deistler, Jan-Matthis Lueckmann, Conor Durkan, Pedro J.
Gonçalves, David S. Greenberg. and Jakob H. Macke. sbi: A toolkit for simulation-based inference. Journal of Open Source Software, 5(52):2505, 2020. doi 10.2000
Veitch et al. Parameter estimation for compact binaries with ground-based gravitational- wave observations using the LALInference software library. Phys. Rev, D91(4):042003, 2015. 1017
13

---

Published as a conference paper at ICLR 2022
A DERIVATIONS
A.
EQUIVARIANCE RELATIONS
Consider a system with an exact equivariance under a joint transformation of parameters and observations
- 0' = g0, + x' = =Tgz
(15)
(16)
An invariant prior fulfills the relation
p = p(6') det Jgl+
(17)
where the Jacobian J, arises from the change of variables rule for probability distributions. A similar relation holds for an equivariant likelihood,
p(= p(='10') |det Jrl-
(18)
An invariant prior and an equivariant likelihood further imply for the evidence p(z)
p(x) / 0)p()d0 = p(x'100 ((0''10)) Idet Jrl |det Jوl d0 = p(x') det Jrl
(19)
Combining an invariant prior with an equivariant likelihood thus leads to the equivariance relation
p(0)2) p(=10)p(() p(a'10") /det Jr p(0') |dct Jyl
p(z)
p(x') det Jrl
(20)
r(0' (x') Idet Jgl
for the posterior, where used Bayes' theorem and equations (17
18 and (19)
A.2 EQUIVARIANCE OF p))
Here derive that an equivariant posterior p(0|z) remains equivariant if the distribution is also
conditioned on the proxy I. as used in equation (7). With the definition p(ğl0) = ι. ((g")=1) () from
section 3.3
p () is equivariant under joint application of E G to 9 and 0,
p
((g")-1
(
((hg²) -1 hg) (hğ) ho).
(21)
where for the last equality we used gh = hg" This implies, that p(ğ)zr is equivariant under joint application of h and Th.
x( = p(010,x))((01 do / p(hğ (1)h)())))((((TTaat) /de Jul d0
(22)
= (hh))/Tiz).
in the second step we used p(g]0,x)) = (() From these relations, the equivariance relation used in equation 7p follows,
p00 p(0,ğ(z))
p(hğ (h0pp()))TTa)))
|det Jul
p(0(x))
p(Glz)
p(hg) Tha)
(23)
= pphhhO0Taa hg) Idet Jal-
A.3 EXACT EOUIVARIANCE OF INFERRED POSTERIOR
Consider posterior that is exactly equivariant under G,
p00Iz) p(g0 T det Jgls
Vg E G.
(24)
We here show that the posterior estimated using GNPE is equivariant under G by construction. This holds regardless of whether (0'(2z') has fully converged to p(0' (a').
14

---

Published as a conference paper at ICLR 2022
With GNPE, the equivariant posterior p(0(z.) for observation T. is inferred by alternately sampling
g( p(0|zo. 00 1
g(+) (( ge()) q(6" -( (ighs-))-
(25)
g(t) ~0 B())
g(0 g0 C, c~ k(c),
(26)
see also equation 9 Now consider a different observation To = T, that is obtained by altering the pose of 2o with Tu. where h is arbitrary element of the equivariance group G. Applying the joint transformation
9 + hg.
(27)
To + Thoo
(28)
in 25 leaves l' invariant,
q ((na) [2 (T,) q(f')TT-1Z0).
(29)
We thus find that 0 in 25 transforms equivariantly under joint application of and 28
0 == =90 (hg))' h(g0") h0.
(30)
Conversely, applying
0 - h0
(31)
in 226 transforms 9 by
g g=g'e (h0) c hgg hg.
(32)
The 0 samples (obtained by marginalizing over ()) thus transform 0 - h0 under vo - Th which 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 H, (initialized with g( = gz.
and T. (initialized with g) = hg,,). ). The Gibbs sampler will yield parameters samples (0, )= for Fo. and the exact same samples (N for Ty up to the global transformation by h. The reason is that the density estimator q(0"(x') (x') is queried with the same x' for both observations x,, and T. in each iteration 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(0|=', g) to obtain samples from the joint distribution p p00, 9 (lr). This is done by iterative sampling as described in section 3 Here we derive equation which states how a distribution (0)lz) ) is updated by single GNPE iteration.
Given a distribution (0(x)) of 0 samples in iteration j. we infer samples for the pose proxy 9 for the next iteration by (i) extracting the pose g from 0 (this essentially involves marginalizing over all non pose related parameters) and (ii) blurring the pose 9" with the kernel K, corresponding to group convolution A.I We denote this combination of marginalization and group convolution with the "*" symbol,
0)-1(01=)= = / do do
(0).
(33)
For given proxy sample g. a (perfectly trained) neural density estimator infers with
g2
<< <x((g))~ p(0lz,0)p(( p(0)x)
p(0 9)
(34)
p(]lx)
p(ğ(x)
(-'(-1x) xx) (g)
A. We define group convolution as (A ()) dg A(g) 8), which is the natural extension of standard convolution.
15

---

Published as a conference paper at ICLR 2022
where used p(ğ xx(("))10)) Combining (33) and 34 samples reads
the updated distribution over
(011)) / dg p(0 9) (011)
<((g")>'0) 'g)
dg (09(1Z)) * ()) (( (1(a) ( * k))
= -(01z) / do0 (1)(
(35)
(12((1)) #x) x((g")
(-)(11)) &
= =(0]x=)
R (-)
(g).
(+(+)=)*
Here, R denotes the reflected kernel, /c (g) x(g 1) Vg. Since we choose symmetric kernel in practice, use x x in
In this notation, the initialization of the pose gd in iteration with Bmit simply means setting Qo(-Iz) * (-x)t-122)
B GNPE FOR SIMPLE GAUSSIAN LIKELIHOOD AND PRIOR
Consider simple forward model т æ with a given prior
p(T) = (-5,, [r]
(36)
and likelihood
p(=/r) = (r, [x]
(37)
where the normal distribution is defined by
[a p)
exp
2x
(u, o² [x]]
(38)
V2TO
The evidence can be computed from the prior (36) and likelihood
37 and reads
p(x)= / drp(T)p(x(r) = / M
-
(39)
The posterior is then given via Bayes' theorem and reads
p(T[z) p(xlr)p(T) W(r, () N
(
xx-5
1/2 [T].
(40)
p(x)
W 5, 22)
B. EQUIVARIANCES
The likelihood
37 is equivariant under G, i.e., the joint transformation
- gT 7 AT,
(41)
This follows directly from (3) and
If the prior invariant, then this equivariance would be However, the prior is not invariant. It turns out that the posterior is still equivariant, but transforms under different representation than it does for the equivariance of the likelihood. Specifically, the posterior 440) is equivariant under joint transformation
38
inherited by the posterior, see App. A.
- gr + AT, 1px== -
(42)
which again directly follows from 40 and (38). Importantly, T3 * T i.e., the representation under which transforms is different for the equivariance of the likelihood and the posterior. For GNPE, the relevant equivariance is that of the posterior, i.c. the set of transformations (42). see also equation ④
The equivariance relation of the posterior thus reads
p(T(x)) p(gT| (75xx) x)) det. Jzl-
Vg E G.
(43)
16

---

Published as a conference paper at ICLR 2022
p(r(2)) for x =
ground truth posterior
GNPE samples
0.1
p(rle)
0.3
0.2
0.1
0.0
Figure B.I: Posterior p(r|= =3) (blue) and the corresponding inferred GNPE samples (orange).
B.2 GNPE
We choose as the pose, which we aim to standardize with GNPE. We define the corresponding proxy as
= c~ x(c) = (0 1)
(44)
We can use GNPE to incorporate the exact equivariance of the posterior by construction. To that end define
7' = -
(45)
x' = T g(=t)
T r F.
We then train a neural density estimator to estimate p(T'(z²). This distribution of the same form as p(r (r) and simply given by
x' f [T']
p(1
N
(46)
2
due to the equivariance For GNPE,
We here assume a neural density estimator that estimates 466 perfectly.
Initialize (1) 0:
2. Sample ((1) by (1) == 7(1) c, e~ No 1)]e). and compute 7' and x' via 45
3. Sample (2) by (2) = (2) + +(1, with (2) p(r'(=') = N { 1/2
[T'::
and repeat (2) and (3) multiple times. This constructs a Markov chain. To obtain (approximately independent) posterior samples 7 ~ p(T[z), we truncate to account for burn-in, in, thin the chain and marginalize over We find that the chain indeed converges to the correct posterior (40) see Fig. B.
C TOY EXAMPLE
C.I FORWARD 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 infinitely short pulse. The time evolution of that system is governed
17

---

Published as a conference paper at ICLR 2022
by the differential equation
d²
d
dt (t) + 2,w0
dt 2()) + wāz(t) = 6(t - r),
(47)
where we denotes the undamped angular frequency and / the damping ratio. The solution for the time series x(t) is given by the Green's 's function for the corresponding differential operator and reads
SST
x(t) Bwo(t-T). sin (V1-B1/olt--))
(48)
> T.
v1 - 1-5
This equation describes a deterministic, injective mapping between parameters 0 = (wo: 3,7) and time series observation v,
x = f())
(49)
This implies likelihood p(z 6(x f(0). and thus a point-like posterior. To showcase (G)NPE on this toy problem we introduce stochasticity by setting
l0 80)
(50)
instead. We sample d0 from uncorrelated Gaussian distribution
o w0
80 ~(0,E,,
E=
(51)
with Owo 0.3 Hz. = 0.03 and 0, = 0.3 Due to the injectivity of f, the posterior p(0 reduces to the probability p(50 = fz) 0). With a uniform prior, and neglecting boundary effects, this implies uncorrelated Gaussian posterior p(0|2;) centered around (-) with standard deviations as specified above. We choose this approach over, e.g., adding noise straight to observations keep the problem simple possible, such that the focus remains the comparison of GNPE and NPE. In particular, knowing that the ground truth posteriors are Gaussian, we can use simple Gaussian density estimator.
We choose uniform priors
p(wo) U[3 10) Hz,
- ((022,0)
= =1-5
(52)
The observational data is the discretized time series in the interval [-5, +5] 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 convolutional embedding network with kernel sizes [5,5,5], stride 1, (2.12 channels, average pooling with kernel size 7 and stride 7, and ReLU activation functions. GNPE, we use the same architecture as for NPE for both, (0'(xx²) and Gmitt (mii(T (T(z). For further hyperparameters, we use the defaults of the sbi package [Tejero-Cantere et al. 2020).
an embedding network. For NPE-CNN: we use three-layer
C.3 RESULTS
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.
D GRAVITATIONAL WAVE PARAMETER INFERENCE
D.1 FORWARD MODEL AND AMORTIZATION
The forward model mapping binary black hole parameters 0 (Tab. D.II to simulated measurements x in the detectors consists of stages. Firstly, the waveform polarizations h(()) for given parameters
18

---

Published as a conference paper at ICLR 2022
Table D.I: Priors for the astrophysical binary black hole parameters used train the inference network. Priors are uniform over the specified range unless indicated otherwise. We train networks with different distance ranges for the two observing runs OI and O2 due the different detector sensitivities. At inference time, cosmological distance prior is imposed by reweighting samples according to their distance.
Description
Parameter
Prior
component masses spin magnitudes spin angles
M11 M2
I0, 80) Mo. me 2 1/2
01, 02 01, 82 112, ØJL Lc
0. 0.88
standard as in Farr et al.
(014)
time of coalescence luminosity distance
[-0.1, 0.1]
dr.
01, 2 detectors: [100, 2000] Mpc
02, detectors: 100, 2000]
Mpc and [100, 6000] Mpc
02, 3 detectors: [100, 1000]
Mpc
reference phase inclination
d
[0, 2T]]
0, π uniform in sine
polarization
+b
0, π
sky position a, 3
uniform over sky
8 are computed with the waveform model IMRPhenomPy2 (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 Jr 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. drifts over the duration of an observing run. To fully amortize the computational cost, we use variety of different PSDs S, 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 the time of the event, see see Dax et (202I) for details. Since this conditioning on S, has no effect on the GNPE algorithm outlined in this work, we keep it implicit in all equations.
D.2 NETWORK 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 flow, that takes this feature vector input. For each detector, the input the embedding network consists of the complex-valuee frequency domain strain in the range [20 Hz, 1024 Hz] with a resolution of 0.125 Hz, and PSD information (1046 S.) with the same binning. This results to a total of (3 2,,3)) real input bins per detector. The first module of the embedding network consists of 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 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-ite it is for instance not compatible with a convolutional neural network. Following this compression layer, we use a series of 24 culle-coneected residual blocks with two layers each 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 Gee is done after the embedding network, by concatenating fie. to the embedded feature vector.
Following this, we use neural spline flow (Durkan et al. 2019) with rational-quadrati spline coupling transforms as density estimator. We use such transforms, each of which is associated with 5 two-layer residual blocks with hidden dimension 5 512. In total, the inference network has 348 hidden layers and 1.31 10th (for two detectors) or 1.42 10th (for three detectors) learnable parameters.
We train the inference network with data set of 5 106 waveforms with parameters 0 sampled from the priors specified in table D.I and reserve 2% of the data for validation. We pretrain the network with learning rate of 3 10 for 300 epochs with fixed PSD, and finetune for another 150 epochs with learning rate of 3- 10 with varying PSDs. With batch size 4,096, training takes 16- 18 days on
NVIDIA Tesla V100 GPU.
19

---

Published as a conference paper at ICLR 2022
0.70
GW150914
GW170809
GW151012
GW 17081814
GW170101
GW CW17080 170818
0.65
GW170729
GW170823
c2st
0.60
0.55
0.50
NPE
NPE (chained)
GNPE (wide k))
GNPE (narrow x))
Figure D.I: c2st scores quantifying the deviation between the inferred posteriors and the MCMC reference. This is an extended version of Fig.
Figure D.2: Comparison of esti- mated posteriors against LALIN- FERENCE MCMC for eight GW events, quantified by the mean squared error (MSE) the sam- ple means. Before computing the means, we normalize each dimen- sion such that the prior has stan- dard deviation of 1. Ø indicates the average across all eight events. GNPE with narrow kernel con- sistently outperforms the baselines. which is in accordance with Fig.
0.05
GW150014
GW170800
0.04
GW170184
GW1708I1
0.03
MSE
0.02
Ø 0.017
25 0.013
Ø 0.005
05 0.002
0.00
NPE
NPE (chained)
GNPE
GNPE
(wide
(narrow x)
D.3 RESULTS
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 c2st score of 10,000 samples for inferred and target posterior. Fig. displays the mean of the score across 5 different sample realizations, Fig. D.I 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 le (since the reference posteriors generated with LAL LALInference do not contain that variable) and the sky position parameters a and 8 (since the NPE baseline with chain rule decomposition infers these in another basis).
20

---

Published as a conference paper at ICLR 2022
0.05
GW150914
GW1778899
GW151012
CW170814
0.04
GW170101
GW 170818
GW170729
GW170823
0.03
MSE
0.02
0.01
0.00
NPE
NPE (chained)
GNPE (wide k)
GNPE (narrow x))
Figure D.3: MSE between inferred posteriors and MCMC reference. Extended version of Fig. D.2
MCMC
MCMC
NPE
NPE
NPE (chained)
NPE (chained)
GNPE (wide R)
GNPE (wide K)
my [Ma]
GNPE (narrow R)
m22 [Ma] & % 26 & 46
GNPE (narrow R)
02 24
83
mi
[M]]
m2
[Mo]
81
02
[M]]
my
[Ma]
02
Figure D.4: Corner plots for the GW events CW170809 (left) and GW GW170814 (right), plotting
marginals the diagonal and 90% credible regions for the 2D correlations. We display the two black hole masses M, and m2 and two spin parameters 8 and 8g (note that the full posterior is 15-dimensional). This extends Fig. by also displaying the results from chained NPE and GNPE with wide k.
21

---

Published as a conference paper at ICLR 2022
MCMC GNPE
[Ma]
32
88
m2
20
24
01
10
08
0.0
24
02 1.6
0.8
36
[ms
d/
33
tn
86
33
80
1m15
28
37
36
35
0.8
1.6
2.4
0.8
1.
m1
[ [
m2 [M
01
2
th
[ms]
t
[ms]
Figure D.5: Corner plot for the GW event GW170814 plotting 1D marginals on the diagonal and 90% credible regions for the 2D correlations. We display the two black hole masses my and m2 and two spin parameters 0, and 0g that are also shown Fig. 5 We additionally display one of the pose parameters tu and the corresponding proxy fu from the last GNPE iteration. In training, the neural density estimator learned that the true pose t differs by at most ms from the proxy tu that is conditioned on (since chose kernel Knamnw U[-1 ms, 1 ms)" see section 5.2). This explains the strong correlation between tu and f observe. For the same reason, the observed correlations between the εh and the non non-pose parameters (m1 m2, 01, 02) are similar to those between the true pose i and (m,, m2, 0, 02))
22

---

Published as a conference paper at ICLR 2022
MCMC chained NPE, NPE pose chained NPE, oracle pose
26 30
m2
45 20
21
01 oT
08
24
82
16
08
40.0
HH 1
37.5
35.0
32.5
32
t1. (m))
26
36
24
0.8
16
24
0.8
32.5
35.0
40.0
+
m
[M]
M2
[Mo]
9
t
[ss]
tr
[ms]
Figure D.6: Corner plot for GW170814 with the four parameters (m1 m2, 0,, 02) that are also displayed Fig.5) 5 as well as the pose (t f.). We compare chained NPE as described in section 5.3 to an oracle version: for the earlier the pose is inferred using standard NPE (green) while for the latter we take an oracle pose provided by (slow) nested sampling algorithm (teal). We observe, that the result using the oracle pose matches the MCMC reference posterior well, while the other one shows clear deviations. Both versions use the same density estimator for the non-pose parameters ó C This demonstrates that inaccuracies of the chained NPE baselines can be almost entirely attributed to inaccurate inital estimates of the pose. Poor pose estimates can occur since the density estimator trained extract the pose operates
nose-standardized data.
Note: The MCMC reference algorithm LALInference does not provide full pose information since automatically marginalizes over fe. For the oracle pose we thus employ the nested sampling algorithm bilby (Ashton et al. 2019).
23

---

Published as a conference paper at ICLR 2022
32
NPE pose
NPE pose proxy
oracle pose
33
oracle pose proxy
30
iss
ms
30
28
t
i
27
20
24
24
35.0
37.5
40.0
30
32
t ms]
H ms
Figure D.7: Left: Pose parameters н and t for the GW event GW 170814, estimated with the neural density estimator qinit, with standard NPE (green), as well as the "true" pose inferred with bilby (teal). Right: Pose proxies tu and tu for the wide kernel Kwide U ms, 3 ms]"" These are obtained from the pose estimates in the left panel via a convolution with Kwide We observe that the deviation between the oracle and the NPE estimate is substantially smaller for the pose proxy than for the pose itself due to the blurring operation. This leads to better performance of fost-mode GNPE (with widde and only one iteration) compared to chained NPE in section 5.3
24


Reviews:
[{"review_id": "ps-C6Wt72b", "reviewer": "Reviewer_Xr5B", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I found this paper a very enjoyable read, but albeit quite difficult to review for the following reasons. I believe the authors are motivated by extremely interesting problems in physics that leverage non-conventional---atleast in machine learning---datasets like gravitational waves. At best, I can comment from afar that the results look promising and the overall execution is of high quality. \n\nI do have some concerns regarding the technical aspects of the paper though. First, there are a few clarifications that I'd like the authors to address. In all the equivariant literature, including designing equivariant normalizing flows, we are interested in pushing forward a $G$-invariant prior density to a $G$-invariant target. In this paper, the authors instead seek to model a $G$-equivariant density/posterior which is a completely different creature. However, there must be a typo or a misunderstanding on my part as equation 4 is an ***invariant*** distribution and not equivariant as otherwise the group element $g$ must commute and appear in the LHS. This also brings a few points with regards to writing clarity and perhaps mathematical rigor. It is unclear to me what it means for a group element $g$ to act on parameters $\\theta$. For the data $x$ this is not an issue as we take a representation of the group and act on the space, but how does one act on the parameters? Do the authors mean acting on the density $p(\\theta)$? This leads to my central complaint of the paper which is the use of the word pose. Simply put this is vague in the context of equivariance despite the citation to Jaderberg et. al 2015. It seems like the main thrust of defining pose is so that we can find an inverse group element which can then be used to standardize the pose. By the definition of a mathematical group such an element always exist and as a result I did not follow the discussion in section 3.3 that requires a kernel $\\kappa$ over a subset of elements in $G$. Specifically, I do not understand the practical motivation of standardizing the pose to a narrow band. \n\nWith regards to the claims about exact and approximate equivariance achieved by GNPE in comparison to prior work I believe the statement is not the best comparison. To be precise, prior work has focused on guaranteeing equivariance by constructing equivariant maps directly while here equivariance relies on finding a pose standardizing element. It is completely unclear how difficult this could be for many groups of interest and as a result I believe we cannot fully generalize this claim, unless the authors provide an efficient algorithm to do so. \n\nI have no specific comments on the experiments but it does seem that GNPE improves over NPE but the data domain is a bit far removed from my area of expertise.\n\n", "summary_of_the_paper": "This paper studies group equivariant neural posterior estimation which seeks to endow conventional NPE method with equivariance of both the data and parameters simultaneously. To test the efficacy of the proposed approach the authors experiment with gravitational wave data and show GNPE achieves considerable performance gains. ", "summary_of_the_review": "Overall the paper is well written and the experiments seem to demonstrate the power of GNPE. There are a few technical questions that the authors could do well to answer which may change my score and opinion of the paper.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "cjNHIsDawHu", "reviewer": "Reviewer_W1NJ", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The major disadvantage of the paper is the writing, which makes the paper difficult to read and evaluate. I can understand the key idea of the GNPE, and the problem investigated by this paper is certainly important. But I cannot check the correctness of the conclusions and algorithm due to lack of necessary assumptions and derivation steps.\n\nSome comments:\n\n1) Where does Eq. (4) come from? I think this is a corollary of several assumptions including the equivariance of prior distribution of theta. If so, please provide all the related assumptions and the proof, which are not trivial for readers.\n\n2) It is unclear why we need \\epsilon and \\hat g. Are they introduced for the case of approximate equivariance. The first two paragraph in Section 3.3 seems to discuss the necessity of introducing them, but it is really hard to understand. Is it possible to make the explanation more clear by using the toy example?\n\n3) What is the definition of p(\\theta|x,\\hat g) in Eq. (5)? How is Eq. (5) proved? In addition, please provide proofs of all equations in Section 3.\n\n4) Is there any way for users to choose the distribution of \\epsilon in applications?\n\n5) In experiments, only NPE and GNPE are considered. It is worth comparing the performance of GNPE versus other flow models with equivariant architecture.", "summary_of_the_paper": "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.\n\nThe 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.", "summary_of_the_review": "It is an interesting paper, and possibly make important contributions. For now, I give the score \"marginally below the acceptance threshold\" due to the poor writing. But if the authors can well explain all the related conclusions and definitions, I can consider to increase the score.\n\nAfter reading the rebuttal and explanations from authors, I think the revised manuscript is much more understandable, and I change the score to \"marginally above the acceptance threshold\".", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "ZJFerTDXyT", "reviewer": "Reviewer_NKhw", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n- The approach is independent of neural architectures and does not necessitate knowledge of exact equivariances.\n- The method seems to be much better than regular NPE in cases where there are known equivariances.\n\nQuestions:\n- I found the writing in Sections 3.4 and 4 quite confusing to follow. I am not sure how changing $g_\\theta$ to the approximate $\\hat{g}$ solves the problem of not having \"pose\" data in images. \n- How is $q_{init}$ trained? Is this a regular NPE trained on the same data used for training invariances? \n- I think Equation 9 needs further discussion in the main body. What is $Q$ and how is it related to lowercase $q$?\n- I think it would be really helpful to show in some toy cases how $q(\\theta | x , \\tau)$ depends on $\\tau$ (or $\\hat{g}$).", "summary_of_the_paper": "The paper proposes to incorporate invariances of posterior estimates via a factored posterior estimation approach rather than the standard architectural interventions. The paper proposes to learn an approximate standardizing transformation over data and then learning posteriors over the approximate standardized data.", "summary_of_the_review": "I am awaiting author responses at this point and will update my review based on their explanations.  ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "TGl41CtbEv6", "reviewer": "Reviewer_tFyT", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Overall I found the problem and approach to be interesting.  The paper is well-written, but I have a number of comments that I list below.\n\nMajor comments:\n- Despite being central to the main application and some of the setup of the paper, it was not clear to me what \"approximate equivariance\" looks like in practice.  For the damped harmonic oscillator toy example, it's clear that the data are exactly jointly equivariant with $\\tau$: if $\\tau$ is shifted by 1, the data are also shifted by 1.  In the approximately equivariant case, I'm not sure what that would look like.  Is it that the distribution of the other parameters depends on $\\tau$ but only weakly?  In the real data example, it's clear what transformation should be applied to the data to put it in the standardized pose for the exact equivariance, but what do the approximately equivariant data transformations look like?\n- The \"Chained NPE\" (equation (12) in Section 5.3) makes sense as a reasonable baseline, but it is unclear to me why it performs so much worse than the wide $\\kappa$ GNPE, and it would be good to include more discussion of this point.  In particular, it seems like since the GNPE with wide $\\kappa$ only performs one Gibbs sampling iteration, it essentially 1) draws the initial pose estimates from a standard NPE density estimator, 2) draws the a full set of parameters according to equation (8), 3) draws a noisy version of the pose from equation (7), and then 4) redraws the full set of parameters from equation (8).  If I understand the chained NPE approach from its minimal description, it is essentially only performing steps 1) and 2).  If this is the case, then it feels like having a wide $\\kappa$ is just \"blurring\" the estimate of the pose relative to the chained NPE approach, in which case the increased performance going from chained NPE to wide $\\kappa$ NPE is attributable just to having a more diffuse estimate of the pose in the latter.  Does this indicated that chained NPE is overly confident in its estimation of the pose?  Could chained NPE be \"fixed\" by learning a more expressive posterior distribution (e.g., mixture of Gaussians instead of a single Gaussian)?  I ask because chained NPE is conceptually much simpler and doesn't require the additional tuning parameter of $\\kappa$, or Gibbs sampling, making chained NPE -- at least in theory -- much more attractive.\n\nMinor Comments:\n- Beyond the standard issues with choosing neural network architecture, this approach introduces two new algorithmic parameters that need to be tuned: 1) the \"noise tolerance\" distribution in the pose estimation, $\\kappa(\\epsilon)$ and 2) the number of Gibbs Sampling steps that are needed to reach convergence.  2) does not seem so problematic in practice (there are various methods of diagnosing the convergence of Markov chains) although it would be good to describe how this was chosen for this paper (e.g., on p. 8 it is stated that \"We find that fast-mode GNPE converges in one iteration, whereas accurate-model requires 30.\" but it is not stated what was used to determine this).  To me 1) seems more problematic since it is (in general) a distribution, possibly over multiple dimensions.  Is it sufficient to only consider uniform distributions, as in the present paper, or would other distributions result in better performance?  How should one set the parameters of such distributions?\n- In the paragraph following equation (8) on p. 5, burn-in and thinning only result in approximately independent samples.  Similarly, It should be noted that equations (7) and (8) only result in approximate samples from the posterior (due to 1) the asymptotics of Gibbs sampling; and 2) any potential mismatch between the variational approximation $q(\\theta | x, \\hat{g})$ and the true posterior $p(\\theta | x, \\hat{g})$ due to e.g., the variational family not being rich enough, an insufficiently flexible inference network that cannot learn the amortized posterior, not enough samples when training the inference network, or failure in the optimization of the neural network).  \n- I find the c2st metric a bit difficult to interpret.  It would also be good to consider more interpretable metrics, like the MSE of the posterior mean for each method, as well as e.g., the coverage of credible intervals.\n- In contrast, Figure 4 is very nice and easy to interpret.  It would be helpful to also visualize these same marginal posteriors for the \"wide $\\kappa$\" and \"NPE (chained)\" methods to see the effect of $\\kappa$ (see above comment).\n\n\nTypo:\n- on p. 9 \"GNPE achieves a scores\" --> \"GNPE achieves a score\"", "summary_of_the_paper": "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.", "summary_of_the_review": "The paper is well-written and presents an interesting method for a methodologically interesting problem with a nice application.  Some of the key aspects of the paper (e.g., what does an approximate equivariance look like in practice? why does the method outperform chained NPE? how should one set the kernel width?) are a bit unclear, but overall the paper is a good contribution.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper studies group equivariant neural posterior estimation which seeks to endow conventional neural posterior estimation method with equivariance of both the data and parameters simultaneously. To test the efficacy of the proposed approach the authors experiment with gravitational wave data and show that the proposed GNPE achieves considerable performance gains.


Strengths:

- The approach is independent of neural architectures and does not necessitate knowledge of exact equivariances.
- The method seems to be much better than regular NPE in cases where there are known equivariances.

Weaknesses:

- the writing of the paper is not clear, which makes the paper difficult to read and evaluate.
- It is hard to check the correctness of the conclusions and algorithm due to lack of necessary assumptions and derivation steps.
- the authors are knowledgeable about the subject but present material in a slightly callous way which prevents a precise understanding of their techniques.


This is a borderline paper with two reviewers in favor of acceptance and two with a slight tendency to reject. The two negative reviewers did not engage in a discussion with the authors or did not complete that discussion, failing to confirm their ratings or provide an update of those ratings. They also do not seem to give strong arguments for rejection. Based on this, I recommend the paper for acceptance. However, I encourage the authors to take into account the reviewers' comments, especially the part on clarity and rigor, to improve the paper for its camera-ready version.
Decision: Accept (Poster)