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: This study contributes with a new simulation-based inference method, Truncated Marginal Neural Ratio Estimation (TMNRE). It builds on a previous method, Neural Ratio Estimation (NRE), by combining NRE (for directly inferring marginals of the posterior distribution) and the sequential constraining of the prior distribution by truncation (targeted to a particular observation). On standard tasks, TMNRE is shown to perform on par with state-of-the-art methods. Furthermore, the sequential truncation of the prior allows for local amortisation and thus makes it easy to perform (local) consistency tests.