Variational methods for simulation-based inference

1 INTRODUCTION . Many domains in science and engineering use numerical simulations to model empirically observed phenomena . These models are designed by domain experts and are built to produce mechanistic insights . However , in many cases , some parameters of the simulator can not be experimentally measured and need to be inferred from data . A principled way to identify parameters that match empirical observations is Bayesian inference . However , for many models of interest , one can only sample from the model by simulating a ( stochastic ) computer program , but explicitly evaluating the likelihood p ( x|θ ) is intractable . Traditional methods to perform Bayesian inference in such simulation-based inference ( SBI ) , also known as likelihood-free inference scenarios , include Approximate Bayesian computation ( ABC ) ( Beaumont et al. , 2002 ) and synthetic likelihood ( SL ) ( Wood , 2010 ) methods . However , these methods generally struggle with high-dimensional data and typically require one to design or learn ( Chen et al. , 2021 ) summary statistics and distance functions . Recently , several methods using neural density ( -ratio ) estimation have emerged . These methods train neural networks to learn the posterior ( SNPE , Papamakarios & Murray , 2016 ; Lueckmann et al. , 2017 ; Greenberg et al. , 2019 ) , the likelihood ( SNLE , Papamakarios et al. , 2019 ; Lueckmann et al. , 2019a ) , or the likelihood-to-evidence ratio ( SNRE , Thomas et al. , 2021 ; Hermans et al. , 2020 ; Durkan et al. , 2020 ; Miller et al. , 2022 ) . To improve the simulation efficiency of these methods , sequential training schemes have been proposed : Initially , parameters are sampled from the prior distribution to train an estimation-network . Subsequently , new samples are drawn adaptively to focus training on specific regions in parameter space , thus allowing the methods to scale to larger models with more parameters . In practice , however , it has remained a challenge to realize the full potential of these sequential schemes : For sequential neural posterior estimation ( SNPE ) techniques , the loss function needs to be adjusted across rounds ( Greenberg et al. , 2019 ) , and it has been reported that this can be problematic if the proposal distribution is very different from prior , and lead to ‘ leakage ’ of probability mass into regions without prior support ( Durkan et al. , 2020 ) . Both sequential neural likelihood ( SNLE ) and likelihood-ratio ( SNRE ) methods require MCMC sampling , which can become prohibitively slow– MCMC sampling is required for each round of simulations , which , for high-dimensional models , can take more time than running the simulations and training the neural density estimator . Our goal is to provide a method which combines the advantages of posterior-targeting methods and those targeting likelihood ( -ratios ) : Posterior targeting methods allow rapid inference by providing a functional approximation to the posterior which can be evaluated without the need to use MCMC sampling . Conversely , a key advantage of likelihood ( -ratio ) targeting methods is their flexibility– learned likelihoods can e.g . be used to integrate information from multiple observations , or can be used without retraining if the prior is changed . In addition , they can be applied with any activelearning scheme without requiring modifications of the loss-function . We achieve this method by combining likelihood ( -ratio ) estimation with variationally learned inference networks using normalizing flows ( Rezende & Mohamed , 2015 ; Papamakarios et al. , 2017 ; Durkan et al. , 2019a ) and sampling importance resampling ( SIR ) ( Rubin , 1988 ) . We name our approach Sequential Neural Variational Inference ( SNVI ) . We will show that our simulation-based inference methods are as accurate as SNLE and SNRE , while being substantially faster at inference as they do not require MCMC sampling . In addition , real-world simulators sometimes produce invalid outputs , e.g . when a simulation fails . We introduce a strategy that allows likelihood ( -ratio ) targeting methods ( such as SNVI ) to deal with such invalid simulation outputs . A recent method termed “ Sequential Neural Posterior and Likelihood Approximation ” ( SNPLA ) also proposed to use variational inference ( VI ) instead of MCMC to speed up inference in likelihoodtargeting methods ( Wiqvist et al. , 2021 ) . While this proposal is related to our approach , their VI objective is based on the reverse Kullback Leibler ( rKL ) divergence for learning the posterior . As we also show on benchmark tasks , this leads to mode-seeking behaviour which can limit its performance . In contrast , we show how this limitation can be overcome through modifying the variational objective in combination with using SIR for adjusting posteriors . After an introduction on neural network-based simulation-based inference ( SBI ) and variational inference ( Sec . 2 ) , we present our method , Sequential Neural Variational Inference ( SNVI ) ( Sec . 3 ) . In Sec . 4.2 , we empirically show that SNVI is significantly faster than state-of-the-art SBI methods while achieving similar accuracy on benchmark tasks . In Sec . 4.3 , we demonstrate that SNVI is scalable , and that it is robust to invalid simulation outputs : We obtain the posterior distribution of a complex neuroscience model with one order of magnitude fewer simulations than previous methods . 2 BACKGROUND . 2.1 SIMULATION-BASED INFERENCE . Simulation-based inference ( SBI ) aims to perform Bayesian inference on statistical models for which the likelihood function is only implicitly defined through a stochastic simulator . Given a prior p ( θ ) and a simulator which implicitly defines the likelihood p ( x|θ ) , the goal is to identify the posterior distribution p ( θ|xo ) for an observation xo . The simulator is considered to be ‘ black-box ’ , i.e . one can not evaluate p ( x|θ ) and does not have access to the internal states of the simulator , but only to its inputs θ and its outputs x . We focus on improving likelihood-estimation ( SNLE ) and likelihood-ratio-estimation ( SNRE ) methods . SNLE trains a deep neural density estimator ` ψ ( x|θ ) by minimizing the forward Kullback-Leibler divergence ( fKL ) between ` ψ ( x|θ ) and p ( x|θ ) using samples ( x , θ ) ∼ p̃ ( x , θ ) = p ( x|θ ) p̃ ( θ ) from the simulator with L ( ψ ) = − 1N ∑N i=1 log ` ψ ( xi|θi ) . Here , ` ψ ( x|θ ) is a conditional density estimator learning the conditional density p ( x|θ ) from ( θ , x ) pairs , ψ are its learnable parameters , and p̃ ( θ ) is the proposal distribution from which the parameters θ are drawn ( given by e.g . a previous estimate of the posterior or by an active learning scheme , Papamakarios et al. , 2019 ; Lueckmann et al. , 2019a ) . Analogously , SNRE uses a discriminator , e.g . a deep logistic regression network , to estimate the density ratio r ( x , θ ) = p̃ ( x , θ ) p̃ ( x ) p̃ ( θ ) = p ( x|θ ) p̃ ( x ) . ( Hermans et al. , 2020 ; Durkan et al. , 2020 ) , If the proposal is given by the prior , then one can recover the exact posterior density , otherwise the posterior can be recovered up to a normalizing constant ( Durkan et al. , 2020 ) . Once the likelihood ( or likelihood-ratio ) has been learned , the posterior can be sampled with MCMC . In sequential schemes , the proposal p̃ is updated each round using the current estimate of the posterior – thus , computationally expensive MCMC sampling needs to be run in each round . 2.2 VARIATIONAL INFERENCE . We use variational inference ( VI ) to estimate the posterior distribution . VI formulates an optimization problem over a class of tractable distributions Q to find parameters φ∗ such that qφ∗ ∈ Q is closest to the true posterior p ( θ|xo ) according to some divergence D ( Blei et al. , 2017 ) . Formally , φ∗ = argmin φ D ( qφ ( θ ) ||p ( θ|xo ) ) with qφ∗ ( θ ) = p ( θ|xo ) ⇐⇒ D ( qφ∗ ( θ ) ||p ( θ|xo ) ) = 0 . Recent work has introduced normalizing flows as a variational family for VI ( Ranganath et al. , 2014 ; Agrawal et al. , 2020 ; Rezende & Mohamed , 2015 ) . Normalizing flows define a distribution qφ ( θ ) by learning a bijection Tφ which transforms a simpler distribution into a complex distribution p ( θ|xo ) . Normalizing flows provide a highly flexible variational family , while at the same time allowing low variance gradient estimation of an expectation by the reparameterization trick , i.e . ∇φEθ∼qφ [ f ( θ ) ] = Eθ0∼q0 [ ∇φf ( Tφ ( θ0 ) ) ] with θ = Tφ ( θ0 ) ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ; Rezende & Mohamed , 2015 ) . 3 SEQUENTIAL NEURAL VARIATIONAL INFERENCE ( SNVI ) . 3.1 KEY INGREDIENTS . We propose a framework to use variational inference ( VI ) for simulation-based inference . Our method consists of three parts : A learnable likelihood ( or likelihood-ratio ) model , a posterior model ( typically parameterized as a normalizing flow ) to be learned with VI , and sampling importance resampling ( SIR ) ( Rubin , 1988 ) to refine the accuracy of the posterior ( Fig . 1 ) . The likelihood ( ratio ) model ` ψ ( x|θ ) learns to approximate the likelihood p ( x|θ ) or the likelihood-ratio p ( x|θ ) p ( x ) from pairs of parameters and simulation outputs ( θ , x ) . We use the term SNLVI to refer to SNVI with likelihoods , and SNRVI with likelihood-ratios . After a likelihood ( -ratio ) model has been trained , the posterior model qφ ( θ ) is trained with variational inference using normalizing flows . Finally , SIR is used to correct potential inaccuracies in the posterior qφ ( θ ) – as we will show below , the SIR step leads to empirical improvements at modest computational overhead . To refine the likelihood ( -ratio ) model and the posterior , the procedure can be repeated across several ‘ rounds ’ . We opt to sample the parameters θ from the previous posterior estimate qφ ( θ ) , but other strategies for active learning ( e.g . Lueckmann et al. , 2019b ) could be plugged into SNVI . The algorithm is summarized in Alg . 1 . We will now describe three variational objectives that can be used with SNVI , the SIR procedure to refine the posterior , and a strategy for dealing with invalid simulation outputs . 3.2 VARIATIONAL OBJECTIVES FOR SBI . Because of the expressiveness of normalizing flows , the true posterior can likely be approximated well by a member of the variational family ( Papamakarios et al. , 2021 ) . Thus , the quality of the Algorithm 1 : SNVI 1 Inputs : prior p ( θ ) , observation xo , divergence D , simulations per round N , number of rounds R , selection strategy S. 2 Outputs : Approximate likelihood ` ψ and variational posterior qφ . 3 Initialize : Proposal p̃ ( θ ) = p ( θ ) , simulation dataset X = { } 4 for r ∈ [ 1 , ... , R ] do 5 for i ∈ [ 1 , ... , N ] do 6 θi = S ( p̃ , ` φ , p ) ; // sample θi ∼ p̃ ( θ ) 7 simulate xi ∼ p ( x|θi ) ; // run the simulator on θi 8 add ( θi , xi ) to X 9 end 10 ( re- ) train ` ψ ; ψ∗ = argminψ − 1N ∑ ( xi , θi ) ∈X log ` ψ ( xi|θi ) ; // or SNRE loss 11 ( re- ) train qφ ; φ∗ = argminφD ( qφ ( θ ) ||p ( θ|xo ) ) with p ( θ|xo ) ∝ p ( xo|θ ) p ( θ ) ≈ ` ψ∗ ( xo|θ ) p ( θ ) 12 p̃ ( θ ) = qφ ( θ ) 13 end variational approximation is strongly linked to the ability to achieve the best possible approximation through optimization , which in turn depends on the choice of variational objective D. Using the reverse Kullback-Leibler Divergence ( rKL ) as proposed by Wiqvist et al . ( 2021 ) can give rise to mode-seeking behaviour and qφ might not cover all regions of the posterior ( Bishop , 2006 ; Blei et al. , 2017 ) . As a complementary approach , we suggest and evaluate three alternative variational objectives that induce a mass-covering behaviour and posit that this strategy will be particularly important in sequential schemes . 1 . Forward KL divergence ( fKL ) In contrast to the reverse KL ( rKL ) , the forward KullbackLeibler divergence ( fKL ) is mass-covering ( Bishop , 2006 ) . Wan et al . ( 2020 ) minimize the following upper bound to the evidence , which implicitly minimizes the fKL : L ( φ ) = Eθ∼qφ [ w ( θ ) log ( w ( θ ) ) ] with w ( θ ) = p ( xo , θ ) /qφ ( θ ) . This expression is hard to estimate with samples : If qφ ( θ ) is different from p ( xo , θ ) then w ( θ ) ≈ 0 for most θ ∼ qφ ( θ ) , thus∇φL ( φ ) ≈ 0 , which would prevent learning ( see Appendix Sec . A.3 ) . To overcome this problem , we rewrite the fKL using self-normalized importance sampling ( Jerfel et al. , 2021 ) . Let θ1 , . . . , θN ∼ π be samples from an arbitrary proposal distribution π . We then minimize the loss : LfKL ( φ ) = DKL ( p||qφ ) ≈ N∑ i=1 w ( θi ) ∑N j=1 w ( θj ) log ( p ( xo , θi ) qφ ( θi ) ) where w ( θ ) = p ( xo , θ ) /π ( θ ) . As a self-normalized importance sampling scheme , this estimate is biased , but the bias vanishes at rateO ( 1/N ) ( Hesterberg , 2003 ) . In our experiments , we use π = qφ , which provides a good proposal when qφ is close to p ( Chatterjee & Diaconis , 2018 ) . Even though qφ will differ from p initially , sufficient gradient information is available to drive qφ towards p , as we demonstrate in Appendix Sec . A.3 . 2 . Importance weighted ELBO The importance weighted ELBO ( IW-ELBO ) introduced by Burda et al . ( 2016 ) uses the importance-weighted gradient of the evidence lower bound ( ELBO ) . It minimizes the KL divergence between the self-normalized importance sampling distribution of qφ and the posterior and thus provides a good proposal for sampling importance resampling ( Cremer et al. , 2017 ; Domke & Sheldon , 2018 ; Ranganath et al. , 2014 ) . It can be formulated as L ( K ) IW ( φ ) = Eθ1 , ... , θk∼qφ [ log 1 K K∑ k=1 p ( xo , θk ) qφ ( θk ) ] . To avoid a low SNR of the gradient estimator ( Rainforth et al. , 2018 ) , we use the ‘ Sticking the Landing ‘ ( STL ) estimator introduced by Roeder et al . ( 2017 ) . 3 . Rényi α-divergences Rényi α-divergences are a divergence family with a hyperparameter α which allows to tune the mass-covering ( or mode-seeking ) behaviour of the algorithm . For α → 1 , the divergence approaches the rKL . For α < 1 , the divergence becomes more mass-covering , for α > 1 more mode-seeking . We use α = 0.1 in our experiments . A Rényi variational bound was established by Li & Turner ( 2016 ) and is given by Lα ( φ ) = 1 1− α log ( Eθ∼qφ [ ( p ( xo , θ ) qφ ( θ ) ) 1−α ] ) For α = 0 , Lα is a single sample Monte Carlo estimate of the IW-ELBO ( when using K samples to estimate the expectation in Lα ( φ ) ) and thus also suffers from a low SNR as α → 0 ( Rainforth et al. , 2018 ; Li & Turner , 2016 ) . Just as for the IW-ELBO , we alleviate this issue by combining the α-divergences with the STL estimator .

Summary: This paper presents SNVI, a simulation-based inference based likelihood (ratio) estimation and variational inference. The main contribution is to sidestep the slow MCMC sampling step of SNLE (SNRE) by learning the posterior distribution with variational inference. This work is similar to SNPLA but replaces the rKL variational objective with three alternative objectives that induce a mass-covering behaviour and are, therefore, better suited when using the approximate posterior distributions as proposal distributions in the sequential regime. Experiments show convincing results.