Truncated Marginal Neural Ratio Estimation

1 Introduction . Parametric stochastic simulators are ubiquitous in science [ 1 , 2 , 3 ] and using them to solve the Bayesian inverse problem is of general interest . Likelihood-based methods like Markov chain Monte Carlo ( MCMC ) [ 4 , 5 ] or nested sampling [ 6 ] are applicable when the likelihood is tractable . It is equally common that the likelihood is only implicitly defined by the simulator or is inefficient to compute . For this so-called likelihood-free or simulation-based inference , the traditional approach is Approximate Bayesian Computation ( ABC ) [ 7 , 8 ] . See [ 9 ] for a reference . Simulation-based inference ( SBI ) is closely connected to ABC and has been an open research topic since as early as the 1980s [ 10 ] . Deep learning has accelerated progress in the field [ 11 , 12 , 13 , 14 ] . Proposed algorithms that learn the likelihood [ 13 ] or the posterior [ 14 , 15 , 16 ] utilize a density estimator . The likelihood-to-evidence ratio [ 11 ] can be learned via a classification-based technique . Refs . [ 11 ] and [ 14 ] were brought into a unified framework by [ 17 ] . High-fidelity simulators often have many parameters and/or an intractable likelihood function , which can make inference notoriously difficult . Practitioners are usually faced with observational data and an expensive stochastic simulator , without access to the ground truth posterior . They want a testably accurate posterior estimate without extreme simulation expense . With existing methods , the 1Implementation of experiments at https : //github.com/bkmi/tmnre/ . Ready-to-use implementation of underlying algorithm at https : //github.com/undark-lab/swyft/ . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . practitioner must choose between increased accuracy per simulation ( so-called sequential methods [ 13 , 18 ] ) or efficient empirical testability ( so-called amortized methods [ 11 ] ) . We provide a method which offers both simultaneously with a balance that can be tuned by a hyperparameter . Three attributes contribute to this goal : Targeted inference . Focusing simulations on the parameter regions that are most relevant for the inference problem and target data is more efficient . This is particularly true when most posterior density is concentrated compared to the prior ’ s density . Marginal posteriors instead of the joint . Scientific insight is often based on a low dimensional marginalization of the posterior with nuisance parameters removed [ 19 ] . The additional information of the full joint posterior might not be worth the additional cost afforded . Targeting marginals directly , by estimating only the marginal for the parameters of interest , is simpler and sufficient for many scientific , parameter estimation , and bounding purposes . Consistency checks through local amortization . Practitioners are interested in testing the quality of inference methods [ 20 , 21 , 22 ] . One such test is to compare the empirical and nominal contained mass of estimated credible regions [ 23 ] . Amortized methods learn the posterior for any data , generated by any parameter , facilitating empirical study of the nominal credible regions on fabricated data . Still , learning an amortized posterior is excessive if only a small subset of parameters are consistent with a target observation . We propose the concept of local amortization to learn the posterior on said subset , combining simulator efficiency of targeted inference with the testability of amortization . Both are critical components for enabling trustworthy scientific results . Our contribution . We propose an algorithm that simultaneously achieves all three of the above aspects : Truncated Marginal Neural Ratio Estimation ( TMNRE ) . It approximates the marginal likelihood-to-evidence ratio in a sequence of rounds and shines when the joint posterior is prohibitively costly . As a basis , we adopt likelihood-to-evidence ratio estimation proposed in [ 11 ] , although our truncation scheme is applicable to other neural simulation-based inference methods which estimate the posterior or likelihood [ 13 , 17 ] . Our iterative scheme is loosely inspired by likelihood-based nested sampling [ 6 , 24 , 25 ] since we generate training data drawn from a nested sequence of truncated priors in multiple rounds . Our algorithm ( a ) preferentially generates simulations in relevant regions of the parameter space , ( b ) allows estimation of all marginals of interest simultaneously and in parallel from the same training data , and ( c ) yields posteriors that are locally amortized in a constrained region around the posterior , enabling empirical self-consistency test of the inference results . Related work . In Table 1 , we compare the properties and features of a selection of deep-learning based simulation-based inference methods that are directly relevant for our work . Sampling from regions of highest probability density is baked into most likelihood-based methods [ 4 , 5 , 6 , 24 , 25 , 26 , 27 , 28 ] . Amortization is generally not available with these methods because they sample from a particular posterior . Approximate Bayesian Computation ( ABC ) is a rejection sampling technique where proposed samples from the generative model are accepted based on a user defined distance criterion comparing generated data to the observation of interest . Two important methods include REJ-ABC [ 9 ] where the proposal distribution is simply the parameter prior and SMC-ABC [ 29 , 30 ] where the proposal is iteratively refined . Blum and François [ 31 ] introduce an ABC distance criterion weighting mechanism to tune the posterior sampler as well as a proposal prior which draws from a truncated region of true prior . It estimates the support of previously-accepted samples via support vector machines [ 32 ] and samples from this region with rejection . Likelihood-free inference can be cast as a conditional density estimation problem targeting either the posterior directly [ 14 , 15 , 16 , 17 ] or the likelihood [ 13 , 33 ] . This technique requires a density estimator , normally implemented as a mixture density network [ 34 ] or a normalizing flow [ 35 , 36 ] . Neural Likelihood Estimation ( NLE ) performs well on benchmark tasks but must learn a density representation of the data in an unsupervised setting–difficult for complex data . Modern variants of Neural Posterior Estimation ( NPE ) [ 17 ] have become effective enough to offer an alternative estimation method for scientific practitioners [ 37 ] . Amortized Approximate Ratio Estimators / Neural Ratio Estimation ( NRE ) : Binary classification allows estimation of the likelihood ratio between two hypotheses [ 12 , 38 , 39 , 40 , 41 , 42 ] and was most famously applied to Generative Adversarial Networks [ 43 ] . Ref . [ 11 ] noted that naive application in the likelihood-free setting was unsatisfactory because the mathematically arbitrary choice of reference hypothesis significantly affected empirical MCMC results . Comparing likelihoods from jointly drawn ( x , θ ) ∼ p ( x , θ ) and marginally drawn ( x , θ ) ∼ p ( x ) p ( θ ) samples , where x and θ refer respectively to simulated data and simulator parameters , addresses the issue . Ref . [ 17 ] cast NRE and NPE in a unifying framework by adapting the loss function to contrast several possible hypothetical parameters . In this paper we refer to the algorithm described in [ 11 ] as NRE or NRE_A while the likelihood ratio algorithm described in [ 17 ] is referred to as NRE_B . Directly estimating the marginal posterior distribution has been mentioned [ 11 ] and explored [ 44 , 45 ] . Moment networks [ 45 ] produce the ( central ) moments of the posterior distribution , without calculating the density explicitly , via a hierarchy of neural networks trained on a regression problem . Ref . [ 45 ] also introduced a method which learns marginal posteriors with normalizing flows but does not address targeted inference or testability of estimated posteriors . Sequential Methods : The neural likelihood-free methods generally offer a so-called sequential formulation that targets the posterior of a particular observation xo [ 11 , 13 , 17 ] . Rather than drawing samples from the prior , the simulation budget is divided between rounds and the previous round ’ s posterior is used as the new proposal distribution for the next round . This method increases simulation efficiency , but does not allow for amortization . Importantly , sequential methods can become highly inefficient when targeting multiple marginal posteriors directly because the previous round ’ s marginal posterior does not update beliefs about the other parameters . A full parameter vector is necessary to run the simulator , thus defeating the purpose for all nuisance ( marginalized-over ) parameters . 2 Method . We aim to estimate any marginal posterior of interest using an approximate marginal likelihood-toevidence ratio . Although we normally compute every one- and two-dimensional marginal posterior for visualization purposes , our method is not limited to this restriction . We define the object of study ... Let parametric stochastic simulator g be a nonlinear function that maps a vector of real parameters θ = ( θ1 , . . . , θD ) and a stochastic latent state z ∈ RNz , Nz ∈ N to an observation x = g ( θ , z ) . We consider a factorizable prior p ( θ ) = p ( θ1 ) · · · p ( θD ) over the parameters . The joint posterior is given via Bayes ’ rule as p ( θ | x ) = p ( x | θ ) p ( θ ) /p ( x ) , where p ( x | θ ) is the intractable likelihood ( or implicit distribution [ 10 , 46 ] ) and p ( x ) is the evidence . Our goal is to efficiently compute arbitrary marginal posteriors , p ( ϑ | x ) . Here , ϑ are the parameters of interest , and we denote all other ( nuisance ) parameters by η , such that θ = ( ϑ , η ) . The marginal posterior is obtained from the joint distribution p ( ϑ , η | x ) : = p ( θ | x ) by integrating over all components of η , p ( ϑ | x ) = ∫ p ( ϑ , η | x ) dη = ∫ p ( x | ϑ , η ) p ( η ) dη p ( x ) p ( ϑ ) : = p ( x | ϑ ) p ( x ) p ( ϑ ) . ( 1 ) where we used Bayes ’ rule , prior factorizibility , and defined the marginal likelihood p ( x | ϑ ) . 2.1 Marginal Neural Ratio Estimation ( MNRE ) . This paper considers the set of one- and two-dimensional marginal posteriors and their corresponding parameters of interest . Given parameter vector θ ∈ RD , define the set of all parameters associated with the one dimensional marginal posteriors by Θ1 : = { θ1 , . . . , θD } . We do something similar , up to symmetry , for all two dimensional marginal posteriors Θ2 : = { ( θi , θj ) ∈ R2 ∣∣ i = 1 , . . . , D , j = i+ 1 , . . . , D } . We set our marginals of interest { ϑk } : = Θ1 ∪ Θ2 but in the general case , { ϑk } can be any set of marginals that the practitioner desires . For every ϑk we use NRE [ 11 ] to estimate the corresponding marginal likelihood-to-evidence ratio rk ( x | ϑk ) : = p ( x | ϑk ) p ( x ) = p ( x , ϑk ) p ( x ) p ( ϑk ) = p ( ϑk | x ) p ( ϑk ) . ( 2 ) To this end , we train binary classifiers ρ̂k , φ ( x , ϑk ) to distinguish jointly drawn parameter-simulation pairs ( x , ϑk ) ∼ p ( x , ϑk ) from marginally drawn parameter-simulation pairs ( x , ϑk ) ∼ p ( x ) p ( ϑk ) , where φ represents the parameters of the classifier . A Bayes optimal classifier ρk would recover the density ρk ( x , ϑk ) = p ( x , ϑk ) p ( x , ϑk ) +p ( x ) p ( ϑk ) . Then the ratios of interest can be estimated by r̂k ( x | ϑk ) : = ρ̂k , φ ( x , ϑk ) 1− ρ̂k , φ ( x , ϑk ) ≈ p ( x , ϑk ) p ( x ) p ( ϑk ) = rk ( x | ϑk ) . ( 3 ) We train each ratio estimator r̂k ( x | ϑ ) using Adam [ 47 ] to minimize the binary cross-entropy ( BCE ) ` k = − ∫ [ p ( x | θ ) p ( θ ) ln ρ̂k , φ ( x , ϑk ) + p ( x ) p ( θ ) ln ( 1− ρ̂k , φ ( x , ϑk ) ) ] dx dθ . ( 4 ) In practice , we concatenate x with ϑk as the input to ρ̂k , φ . Since each classifier trains independently , it is trivial to train them all in parallel using the same underlying ( x , θ ) pairs . Practically , we parameterize the classifier by ρ̂k , φ ( x , ϑk ) = σ ◦ fk , φ ( x , ϑk ) , where σ is the logistic sigmoid and fk , φ is a neural network . The connection in Eq . ( 3 ) between the estimated ratio and the classifier implies that log r̂k ( x | ϑk ) = fk , φ ( x , ϑk ) . We call the above technique MNRE . When training data is limited , we found empirically ( see Sec . 3.2 below ) that the MNRE approach typically leads to conservative ( i.e. , not overconfident / not narrower than the ground truth ) likelihoodto-evidence ratio estimates , provided early stopping criteria are used to avoid over-fitting of the classifier . At its core MNRE solves a simple , supervised binary classification task rather than a complex , unsupervised density estimation problem . Classification tasks are generally easier to train [ 12 ] , and can rely on battle-tested network architectures .

Summary: The authors suggest the use of neural ratio estimation to target the marginal posteriors directly, thereby circumventing the more difficult problem of obtaining the full Bayesian posterior. The method comes with other benefits such as local amortization as well as a sequential truncation technique for targeted inference.  The method is validated on several benchmark examples.