Group equivariant neural posterior estimation

1 INTRODUCTION . Bayesian inference provides a means of characterizing a system by comparing models against data . Given a forward model or likelihood p ( x|θ ) for data x described by parameters θ , and a prior p ( θ ) , the Bayesian posterior is proportional to the product , p ( θ|x ) ∝ p ( x|θ ) p ( θ ) . Sampling techniques such as Markov Chain Monte Carlo ( MCMC ) can be used to build up a posterior distribution provided the likelihood and prior can be evaluated . For models with intractable or expensive likelihoods ( as often arise in scientific applications ) simulation-based ( or likelihood-free ) inference methods offer a powerful alternative ( Cranmer et al. , 2020 ) . In particular , neural posterior estimation ( NPE ) ( Papamakarios & Murray , 2016 ) uses expressive conditional density estimators such as normalizing flows ( Rezende & Mohamed , 2015 ; Papamakarios et al. , 2021 ) to build surrogates for the posterior . These are trained using model simulations x ∼ p ( x|θ ) , and allow for rapid sampling for any x ∼ p ( x ) , thereby amortizing training costs across future observations . NPE and other density-estimation methods for simulation-based inference ( Gutmann & Corander , 2016 ; Papamakarios et al. , 2019 ; Hermans et al. , 2020 ) have been reported to be more simulation-efficient ( Lueckmann et al. , 2021 ) than classical likelihood-free methods such as Approximate Bayesian Computation ( Sisson et al. , 2018 ) . Hanford Livingston stan dard ized dire ctio n Virgo Training an inference network for any x ∼ p ( x ) can nevertheless present challenges due to the large number of training samples and powerful networks required . The present study is motivated by the problem of gravitational-wave ( GW ) data analysis . Here the task is to infer properties of astrophysical black-hole mergers based on GW signals observed at the LIGO and Virgo observatories on Earth . Due to the complexity of signal models , it has previously not been possible to train networks to estimate posteriors to the same accuracy as conventional likelihood-based methods ( Veitch et al. , 2015 ; Ashton et al. , 2019 ) . The GW posterior , however , is equivariant1 under an overall change in the time of arrival of the data . It is also approximately equivariant under a joint change in the sky position and ( by triangulation ) individual shifts in the arrival times in each detector ( Fig . 1 ) . If we could constrain these parameters a priori , we could therefore apply time shifts to align the detector data and simplify the inference task for the remaining parameters . More generally , we consider forward models with known equivariances under group transformations applied jointly to data and parameters . Our aim is to exploit this knowledge to standardize the pose of the data2 and simplify analysis . The obvious roadblock is that the pose is contained in the set of parameters θ and is therefore unknown prior to inference . Here we describe group equivariant neural posterior estimation ( GNPE ) , a method to self-consistently infer parameters and standardize the pose . The basic approach is to introduce a proxy for the pose—a blurred version—on which one conditions the posterior . The pose of the data is then transformed based on the proxy , placing it in a band about the standard value , and resulting in an easier inference task . Finally , the joint posterior over θ and the pose proxy can be sampled at inference time using Gibbs sampling . The standard method to incorporating equivariances is to integrate them directly into network architectures , e.g. , to use convolutional networks for translational equivariances . Although these approaches can be highly effective , they impose design constraints on network architectures . For GWs , for example , we use specialized embedding networks to extract signal waveforms from frequencydomain data , as well as expressive normalizing flows to estimate the posterior—neither of which is straightforward to make explicitly equivariant . We also have complex equivariance connections between subsets of parameters and data , including approximate equivariances . The GNPE algorithm is extremely general : it is architecture-independent , it applies whether equivariances are exact or approximate , and it allows for arbitrary equivariance relations between parameters and data . We discuss related work in Sec . 2 and describe the GNPE algorithm in Sec . 3 . In Sec . 4 we apply GNPE to a toy example with exact translational equivariance , showing comparable simulation efficiency to NPE combined with a convolutional network . In Sec . 5 we show that standard NPE does not achieve adequate accuracy for GW parameter inference , even with an essentially unlimited number of simulations . In contrast , GNPE achieves highly accurate posteriors at a computational cost three orders of magnitude lower than bespoke MCMC approaches ( Veitch et al. , 2015 ) . The present paper describes the GNPE method which we developed for GW analysis ( Dax et al. , 2021 ) , and extends it to general equivariance transformations which makes it applicable to a wide range of problems . A detailed description of GW results is presented in Dax et al . ( 2021 ) . 2 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 1In physics , the term “ covariant ” is frequently used instead of “ equivariant ” . 2We adopt the language from computer vision by Jaderberg et al . ( 2015 ) . conflict with design considerations such as data representation and flexibility of the architecture , and imposes constraints such as locality . GNPE achieves complete separation of equivariances from these considerations , requiring only the ability to efficiently transform the pose . Normalizing flows are particularly well suited to NPE , and there has been significant progress in constructing equivariant flows ( Boyda et al. , 2021 ) . However , these studies consider joint transformations of parameters of the base space and sample space—not joint transformation of data and parameters for conditional flows , as we consider here . GNPE enables end-to-end equivariances from data to parameters . Consider , by contrast , a conditional normalizing flow with a convolutional embedding network : the equivariance persists through the embedding network but is broken by the flow . Although this may improve learning , it does not enforce an end-to-end equivariance . This contrasts with an invariance , for which the above would be sufficient . Finally , GNPE can also be applied if the equivariance is only approximate . Several other approaches integrate domain knowledge of the forward model ( Baydin et al. , 2019 ; Brehmer et al. , 2020 ) by considering a “ gray-box ” setting . GNPE allows us to incorporate high-level domain knowledge about approximate equivariances of forward models without requiring access to its implementation or internal states of the simulator . Rather , it can be applied to “ black-box ” code . An alternative approach to incorporate geometrical knowledge into classical likelihood-free inference algorithms ( e.g. , Approximate Bayesian Computation , see ( Sisson et al. , 2018 ) ) is by constructing ( Fearnhead & Prangle , 2012 ) or learning ( Jiang et al. , 2017 ; Chen et al. , 2021 ) equivariant summary statistics s ( x ) , which are used as input to the inference algorithm instead of the raw data x . However , designing equivariant summary statistics ( rather than invariant ones ) can be challenging , and furthermore inference will be biased if the equivariance only holds approximately . Past studies using machine-learning techniques for amortized GW parameter inference ( Gabbard et al. , 2019 ; Chua & Vallisneri , 2020 ; Green & Gair , 2021 ; Delaunoy et al. , 2020 ) all consider simplified problems ( e.g. , only a subset of parameters , a simplified posterior , or a limited treatment of detector noise ) . In contrast , the GNPE-based study in Dax et al . ( 2021 ) is the only one to treat the full amortized parameter inference problem with accuracy matching standard methods . 3 METHODS . 3.1 NEURAL POSTERIOR ESTIMATION . NPE ( Papamakarios & Murray , 2016 ; Greenberg et al. , 2019 ) is a simulation-based inference method that directly targets the posterior . Given a dataset of prior parameter samples θ ( i ) ∼ p ( θ ) and corresponding model simulations x ( i ) ∼ p ( x|θ ( i ) ) , it trains a neural density estimator q ( θ|x ) to estimate p ( θ|x ) . This is achieved by minimizing the loss LNPE = Ep ( θ ) Ep ( x|θ ) [ − log q ( θ|x ) ] ( 1 ) across the dataset of ( θ ( i ) , x ( i ) ) pairs . This maximum likelihood objective leads to recovery of p ( θ|x ) if q ( θ|x ) is sufficiently flexible . Normalizing flows ( Rezende & Mohamed , 2015 ; Durkan et al. , 2019 ) are a particularly expressive class of conditional density estimators commonly used for NPE . NPE amortizes inference : once q ( θ|x ) is trained , inference is very fast for any observed data xo , so training costs are shared across observations . The approach is also extremely flexible , as it treats the forward model as a black box , relying only on prior samples and model simulations . In many situations , however , these data have known structure that one wants to exploit to improve learning . 3.2 EQUIVARIANCES UNDER TRANSFORMATION GROUPS . In this work we describe a generic method to incorporate equivariances under joint transformations of θ and x into NPE . A typical example arises when inferring the position of an object from image data . In this case , if we spatially translate an image x by some offset ~d—effected by relabeling the pixels— then the inferred position θ should also transform by ~d—by addition to the position coordinates θ. Translations are composable and invertible , and there exists a trivial identity translation , so the set of translations has a natural group structure . Our method works for any continuous transformation group , including rotations , dilations , etc. , and in this section we keep the discussion general . For a transformation group G , we denote the action of g ∈ G on parameters and data as θ → gθ , ( 2 ) x→ Tgx . ( 3 ) Here , Tg refers to the group representation under which the data transform ( e.g. , for image translations , the pixel relabeling ) . We adopt the natural convention that G is defined by its action on θ , so we do not introduce an explicit representation on parameters . The posterior distribution p ( θ|x ) is said to be equivariant under G if , when the parameter and data spaces are jointly G-transformed , the posterior is unchanged , i.e. , p ( θ|x ) = p ( gθ|Tgx ) |det Jg| , ∀g ∈ G. ( 4 ) The right-hand side comes from the change-of-variables rule . For translations the Jacobian Jg has unit determinant , but we include it for generality . For NPE , we are concerned with equivariant posteriors , however it is often more natural to think of equivariant forward models ( or likelihoods ) . An equivariant likelihood and an invariant prior together yield an equivariant posterior ( App . A.1 ) . Our goal is to use equivariances to simplify the data—to G-transform x such that θ is taken to a fiducial value . For the image example , this could mean translating the object of interest to the center . In general , θ can also include parameters unchanged under G ( e.g. , the color of the object ) , so we denote the corresponding standardized parameters by θ0 . These are related to θ by a group transformation denoted gθ , such that gθθ0 = θ . We refer to gθ as the “ pose ” of θ , and standardizing the pose means to take it to the group identity element e ∈ G. Applying T ( gθ ) −1 to the data space effectively reduces its dimensionality , making it easier to interpret for a neural network . Although the preceding discussion applies to equivariances that hold exactly , our method in fact generalizes to approximate equivariances . We say that a posterior is approximately equivariant under G if ( 4 ) does not hold , but standardizing the pose nevertheless reduces the effective dimensionality of the dataset . An approximately equivariant posterior can arise if an exact equivariance of the forward model is broken by a non-invariant prior , or if the forward model is itself non-equivariant .

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