Score Modeling for Simulation-based Inference
Abstract
Neural Posterior Estimation methods for simulation-based inference can be illsuited for dealing with posterior distributions obtained by conditioning on multiple observations, as they may require a large number of simulator calls to yield accurate approximations.
Neural Likelihood Estimation methods can naturally handle multiple observations, but require a separate inference step, which may affect their efficiency and performance.
We introduce a new method for simulationbased inference that enjoys the benefits of both approaches.
We propose to model the scores for the posterior distributions induced by individual observations, and introduce a sampling algorithm that combines the learned scores to approximately sample from the target efficiently.
Introduction
Mechanistic simulators have been developed in a wide range of scientific domains [3].
Often these simulators act as a black box: given a set of parameters, the simulator can be sampled, but the distribution over the outputs–the likelihood–cannot be evaluated, rendering typical inference algorithms inapplicable.
Simulation-based inference (SBI) methods provide a way to perform inference with these models [1, 3].
Given a prior over parameters p ( θ ) and a simulator for the likelihood p ( x | θ ) , the goal of SBI is to approximate the posterior p ( θ | x o 1 , . . .
, x on ) for any set of i.i.d. observations { x o 1 , . . .
, x on } .
Most SBI methods work by running the simulator to generate samples x ∼ p ( x | θ ) for different parameters θ , and using the resulting samples to build an approximation of the posterior.
Since many domains involve expensive simulators, recent work has focused on developing algorithms that yield good approximations using a limited budget of simulator calls.
Recent work introduced Neural Posterior Estimation (NPE) methods [20, 17, 8, 2], which use samples ( θ, x 1 , . . .
, x n ) ∼ p ( θ ) (cid:81) nj =1 p ( x j | θ ) to train a conditional neural density estimator q ψ ( θ | x 1 , . . .
, x n ) with parameters ψ , often a normalizing flow [26, 33, 37], via maximum likelihood.
After training, the estimator provides an amortized approximation to p ( θ | x o 1 , . . .
, x on ) for any set of observations { x o 1 , . . .
, x on } of size n .
The drawback of NPE methods is that each training sample requires n simulator calls, so building a training set of size M requires running the simulator nM times.
This may be problematic in scenarios where n is large and calls to the simulator are expensive.
Neural Likelihood Estimation (NLE) methods [21, 39, 15] are a natural alternative for cases where n > 1 .
These methods learn a surrogate likelihood q ψ ( x | θ ) (or a likelihood ratio [22, 4, 10]) using samples ( θ, x )
∼ ˜ p
( θ ) p ( x | θ ) , where ˜ p ( θ ) is a proposal distribution, which in the simplest case can default to the prior p ( θ ) .
Then, given a set of observations { x o 1 , . .
, x on } , inference is carried out on the approximate unnormalized target p ( θ ) (cid:81)
nj q ψ ( x oj | θ ) by standard methods, typically MCMC [21, 15] or variational inference [38, 7].
While these methods can handle arbitrary sets of observations at inference time without re-training the surrogate, they do not approximate the posterior directly, and thus require the inference step to be repeated for each set of observations of interest.
Moreover, their performance depends on the performance of the underlying generic inference methods, which tend to struggle e.g. with multimodal distributions.
Our goal is to develop a method that enjoys the benefits of both types of approaches while avoiding their drawbacks—a method that approximates the posterior directly, is trained on samples ( θ, x ) ∼ p ( θ ) p ( x | θ ) , and is able to naturally handle a varying number of observations at test time.
We propose such an approach that relies on score-based modeling [30, 31, 11, 32].
Simply put, we train a single conditional score network to approximate the score of (diffused versions of) p ( θ | x ) for any x , and propose an algorithm that uses the trained network to approximately sample the posterior p ( θ | x o 1 , . . .
, x on ) for any set of observations { x o 1 , . . .
, x on } .
Our method satisfies the three desiderata outlined above: it directly approximates the posterior, learns form samples ( θ, x ) ∼ p ( θ ) p ( x | θ ) produced with a single call to the simulator, and provides a sampling algorithm that can handle arbitrary sets of observations without re-training.
1.1 Conditional score-based generative modeling
The goal of conditional generative modeling is to learn an approximation of a target distribution p
( θ | c ) for some conditioning variable c given samples ( θ, c ) ∼ p ( θ, c ) , which is exactly the problem SBI methods need to solve.
Methods based on score modeling have shown impressive performance for this task [31, 5, 12, 25, 28, 29].
They define a sequence of conditional densities p 0 ( θ | c ) , . .
. , p T ( θ | c ) by diffusing the target p ( θ | c ) with Gaussian kernels of increasing levels of noise, learn the scores of each density in the sequence using denoising score matching [13, 35], and use Langevin dynamics [27, 36] with the learned scores to approximately sample from the target distribution.
Specifically, for the noise levels
= γ T < γ T − 1 < . . .
< γ 1 < 1 and the corresponding Gaussian diffusion kernels p t ( θ | θ ′ )
=
N ( θ |√ γ t θ ′ , (1 − γ t ) I ) , the sequence of densities is defined [-asp 
-]as p 0 ( θ | c ) =
p ( θ | c ) and p t ( θ | c ) = (cid:90) dθ ′ p ( θ
[Equation]
| c ) p t [Equation] | [Equation]
for t
[Equation] [Equation]
T.
[Equation]
Since γ T = 0 , this sequence gradually bridges between the initial tractable reference N ( θ | 0 , I ) =
p T ( θ ) and the target p ( θ | c ) = p 0 ( θ | c ) .
Score-based methods train a score network s ψ ( θ, t, c ) parameterized by ψ to approximate the scores of these densities, ∇ θ log p t ( θ | c ) .
As only samples from the target are available, this is done via denoising score matching [13, 35], minimizing
[Equation]
Finally, the score network is used to approximately sample the target using annealed Langevin dynamics, as shown in Algorithm 1.
Algorithm 1 Approximate sampling with scores
Input: Score network s ψ ( θ, t, c ) ≈ ∇ θ log p t ( θ | c ) , reference distribution p T ( θ )
Input: Conditioning variable c , number of Langevin steps L , Langevin step sizes δ t θ ∼ p T ( θ ) ▷ Sample reference for t = T − 1 , T − 2 , . .
, 1 do for s = 1 , 2 , . . .
, L do θ ← θ + δ t 2 s ψ ( θ, t, c ) + √ δ t η ts [ η ts ∼ N (0 , I )]
▷ Unadjusted Langevin step return θ
Score-based Neural Posterior Estimation
This section presents our approach for SBI using score modeling.
Our goal is to develop a method that can be trained using parameter/single-observation pairs ( θ, x ) ∼ p ( θ ) p ( x | θ ) , and that can be used at test time to approximate p ( θ | x o 1 , . . .
, x on ) for arbitrary sets of observations { x o 1 , . . .
, x on } with any cardinality n .
As we explain in Appendix A, a naive application of conditional score based modeling fails to satisfy our desiderata.
Instead, we propose an alternative approach based on score modeling, involving different choices for the bridging densities and the reference distribution.
Our method is based on the observation that p ( θ | x 1 , ..., x n ) ∝ p ( θ ) 1 − n (cid:81) nj =1 p ( θ | x j ) (see Appendix A).
Using this factorization, we propose the sequence of densities
[Equation]
where p t ( θ | x j ) is defined in Eq.
(1), taking c = x j .
This construction has four key properties.
First, the distribution for t = 0 recovers the target p ( θ | x 1 , ..., x n ) .
Second, the distribution for t =
T is a tractable Gaussian p T ( θ | x 1 , . . .
, x n )
=
p T ( θ ) =
N ( θ | 0 , 1 n I ) , 2 and thus can be used as a reference for the process.
Third, the score of the resulting densities can be decomposed in terms of the score of the prior (available exactly) and the scores of p t ( θ | x j ) as
[Equation]
And fourth, the scores ∇ θ log p t ( θ | x j ) can all be approximated using a single score network s ψ ( θ, t, x ) trained via denoising score matching using samples ( θ, x ) ∼ p ( θ ) p ( x | θ ) , as explained in
Section 1.1.
After training, given an arbitrary set of observations { x o 1 , . . .
, x on } we can approximately sample the target p ( θ | x o 1 , . . .
, x on ) by running Algorithm 1 with the reference distribution p T ( θ ) =
N ( θ | 0 , 1 n I ) , conditioning variable c = { x o 1 , . .
, x on } , and the approximate score given by
[Equation]
It is straightforward to verify that our approach satisfies our original desiderata: the score network s ψ ( θ, t, x ) is trained using samples ( θ, x ) ∼ p ( θ ) p ( x | θ ) , and the sampling algorithm can be used with any set of observations { x o 1 , . . .
, x on } , as it relies on Langevin dynamics with Eq.
2.1 Alternative sampling approach
The sampling process described in Algorithm 1 requires choosing step-sizes δ t , the number of steps L per noise level, and has complexity O ( LT ) .
This section introduces a different method to approximately sample the target p ( θ | x 1 , . . .
, x n ) , which does not use Langevin dynamics and runs in
O ( T ) steps.
The approach is based on the formulation by Sohl-Dickstein et al.
[30] to use diffusion models to approximately sample from a product of distributions.
The final method involves sampling
T Gaussian transitions with means and variances computed using the learned score network.
We describe the method in Algorithm 2, and give its derivation in Appendix C.
Algorithm
Input s ψ θ
≈ ∇ θ p T
Input
Define α 1 = γ 1 , α t = γ t / γ t − 1 , and β t = 1 − α t , for t = 1 , . .
, T −θ ∼ p T ( θ )
▷ Sample reference for t = T − 1 , T − 2 , . .
, 1 do µ jt
= 1 √ α t ( θ + (1 − α t ) s ψ ( θ, t, x j )) for j = 1 , . . . , n ▷ Compute auxiliary variables σ 2 t = β t
[Equation]
, µ
[Equation]
=
[Equation]
n
[Equation]
+
[Equation]
∇
[Equation]
log [Equation] ( θ ) ▷ Set transition mean and variance θ ∼ N ( θ | µ t , σ 2 t I )
▷ Sample transition return θ
Empirical Evaluation
This section presents an empirical evaluation on two problems commonly used to evaluate SBI methods [16].
One involves a “simulator” consisting of a Gaussian prior and likelihood, p ( θ ) =
N ( θ | 0 , I ) and p ( x | θ ) =
N ( x | θ, Σ) (we set Σ to a diagonal matrix with elements increasing linearly from 0 .
6 to 1 . 4 ), while the other uses a Gaussian prior and a mixture-of-Gaussians likelihood, p ( θ ) =
N ( θ | 0 , I ) and p ( x | θ ) = 0 .
5 N ( x | θ, 2 . 25Σ)
+ 0 .
5 N ( x | θ, 19 Σ) .
In both cases we set θ, x ∈ R 10 .
We compare our approach, called Score NPE, to NPE using a normalizing flow with four Real NVP layers [6] (details in Appendix B).
We compare the methods’ performance when trained on datasets of different sizes, corresponding to different budgets of simulator calls B ∈ { 10 3 , 3 · 10 3 , 10 4 , 3 · 10 4 } .
In all cases we use 20% of the training data as a validation set for early stopping, and train for a maximum of 20k epochs.
We train all methods using Adam [14] with a learning rate of 10 − 4 .
After training we generate a set of observations by drawing θ ∼ p ( θ ) and x o 1 , . . .
, x o 8 ∼ iid p ( x | θ ) , and report the squared MMD [9] between the true posterior and the approximation returned by each method for subsets of { x o 1 , . . .
, x o 8 } of different size.
Figure 1 shows average results over 40 random seeds.
We observe that both methods perform similarly for the simpler Gaussian-Gaussian model, but that Score NPE outperforms the flow baseline for the model with the mixture-of-Gaussians likelihood.
Multimodal posterior.
We also consider a two-dimensional example with a multimodal posterior, with the prior and likelihood given by p ( θ ) =
N ( θ | 0 , I ) and p ( x | θ ) = 0 .
5 N ( x | θ, 0 .
5 I )
+0 .
5 N ( x |− θ, 0 .
5 I ) .
We train each method using a budget of 10 4 simulator calls.
After training we sample θ ∼ p ( θ ) and x o 1 , . . .
, x o 5 ∼ iid p ( x | θ ) , and use each method to generate samples from the approximate posterior obtained by conditioning on subsets of { x o 1 , . .
, x o 5 } of different size.
Results are shown in
Fig.
2, where it can be observed that our method is able to capture both modes well for all subset sizes of observations despite only being trained on parameters/single-observation pairs
Acknowledgments and Disclosure of Funding
We thank Francisco Ruiz for useful comments and suggestions.
A Failure of direct application of conditional score modeling
The target distribution is given by p ( θ | x 1 , . . .
, x n ) .
A direct application of conditional score modeling yields the sequence of densities
[Equation]
It can be seen from the equation above that the score ∇ θ log p t ( θ | x 1 , . . .
, x n ) does not factorize in terms of the single-observation scores ∇ θ log p t ( θ | x j ) , meaning that the corresponding score network would have to be trained using samples ( θ, x 1 , . . .
, x n ) ∼ (cid:81) nj p ( x j | θ ) , obtained by calling the simulator n times for a single sample θ .
As mentioned in Section 1 this is one of the drawbacks of NPE methods that we seek to avoid.
A.1 Derivation of posterior factorization
The factorization for the posterior distribution p ( θ | x 1 , . . .
, x n ) is obtained applying Bayes rule twice:
[Equation]
p ( x j | θ ) (8) ∝ p ( θ ) n (cid:89) j = p ( θ | x j ) p ( θ ) (Bayes rule) (9) = p ( θ
[Equation]
− n n (cid:89) j = p ( θ | x j ) .
B Details for empirical evaluation B.1 Score NPE
Our implementation of the score network s ψ ( θ, t, x ) has three blocks:
• An MLP with 3 hidden layers that takes θ as input and outputs an embedding θ emb , • An MLP with 3 hidden layers that takes x as input and outputs an embedding x emb , • An MLP with 3 hidden layers that takes [ θ emb , x emb , t emb ] as input, where t emb is a positional embedding obtained as described by Vaswani et al.
[34], and outputs the estimate for the score.
(We parameterize the score in terms of the noise variables ϵ [18].)
All MLPs use residual connections throughout.
Running Algorithm 1 to generate samples using the learned score network requires choosing step sizes δ t and the number of Langevin steps L for each noise level γ t .
We use L = 50 and δ t = 0 .
1 − α t √ α t , where α 1 = γ 1 and α t = γ t γ t − 1 for t = 2 , . . .
, T − 1 .
For all our experiments we use T = 400 .
B.2 Flow NPE
We use an implementation of NPE methods based on flows able to handle sets of observations of any size n ∈ { 1 , 2 , . . .
n max } .
The flow can be expressed as q ψ ( θ | x 1 , . . .
, x n , n ) .
Following Chen et al.
[2] and Radev et al.
[23, §2.4], we use an exchangeable neural network to process the observations x 1 , . . .
, x n .
Specifically, we use an MLP with 3 hidden layers to generate an embedding x ej for each observation x j .
We then compute the mean embedding across observations ¯ x e = 1 n (cid:80) nj x ej , which we use as input for the conditional flow.
Finally, we model the flow q ψ ( θ | x 1 , . . .
, x n , n ) =
q ψ ( θ | ¯ x e , n e ) , where n e is an untrained embedding for the number of observations n .
For the flow we use 4 Real
NVP layers [6], each one consisting on MLPs with three hidden layers.
As for the Score NPE method, we use residual connections throughout.
We train the flow via maximum likelihood using samples ( n, θ, x 1 , . . .
, x n ) ∼ Unif( n | min = 1 , max = 10) p ( θ ) (cid:81) nj p ( x j | θ ) .
At test time, this architecture can handle sets of observations of any size n ∈ { 1 , 2 , . .
C Alternative sampling method without unadjusted Langevin dynamics
This section gives the derivation for the sampling method shown in Algorithm 2.
In short, the derivation uses the formulation of score-based methods as diffusions, and has 3 main steps: (1) using the scores of p t ( θ | x ) to compute the Gaussian transition kernels of the corresponding diffusion process [18]; (2) composing n Gaussian transitions corresponding to the n diffusions of p t ( θ | x 1 ) , . . .
, p t ( θ | x n ) (this is based on Sohl-Dickstein et al.
[30]); and (3) adding a correction for the prior term (also based on Sohl-Dickstein et al.
We note that steps 2 and 3 require some approximations.
Despite this, our empirical evaluation indicates that the method works well in practice.
We believe a thorough analysis of these approximations would be useful in understanding when the sampling method from Appendix C can be expected to work.
For clarity, we use [A] to indicate when the approximations are introduced/used.
Connection between score-based methods and diffusion models We begin by noting that scorebased methods can be equivalently formulated as diffusion models, where the mean of Gaussian transitions that act as denoising steps are learned instead of the scores.
Specifically, letting α 1 = γ 1 , α t = γ t / γ t − 1 , and β t = 1 − α t , for t = 1 , . .
, T − 1 , the learned model is given by a sequence of Gaussian transitions p t ( θ t − 1 | θ t , x ) =
N ( θ t − 1 | µ ψ ( θ t , t, x ) , β t ) trained to invert a sequence of noising steps given by q t ( θ t | θ t − 1 ) =
N ( θ t |√ 1 − β t θ t − 1 , β t I ) .
The connection between diffusion models and score-based methods comes from the fact that the optimal means and scores are linearly related [18]
[Equation]
Approximately composing n diffusions To simplify notation, we use a superscript j to indicate distributions that are conditioned on x j (e.g. p jt ( θ t )
=
p t ( θ t | x j ) ).
Assume we have transition kernels p jt ( θ t − 1 | θ t ) that exactly reverse the forward kernels q ( θ t | θ t − 1 ) [A1], meaning that p jt − 1 ( θ t − 1 ) = (cid:82) dθ t p
jt ( θ t ) p
jt ( θ t − 1 | θ t ) , or equivalently p
jt ( θ t ) p
jt ( θ t − 1 | θ t ) =
p jt − 1 ( θ t − 1 ) q t ( θ t | θ t − 1 ) .
Our goal is to find a transition kernel ˜ p t ( θ t − 1 | θ t ) that satisfies
[Equation]
) where ˜ p t ( θ t )
= 1 Z t (cid:81) nj p
jt ( θ t ) .
It is straightforward to verify that the condition from Eq.
(12) can be re-written as
[Equation]
[Equation]
Then, one way to satisfy Eq.
(14) is by setting ˜ p t ( θ t − 1 | θ t ) so that the term in blue above is equal to p 1 t ( θ t − 1 | θ t ) .
That is,
[Equation]
However, the resulting ˜ p t ( θ t − 1 | θ t ) may not be a normalized distribution [30].
Following SohlDickstein et al.
[30], we propose to use the corresponding normalized distribution defined as ˜ p Nt ( θ t − 1 | θ t )
∝ ˜ p t ( θ t − 1 | θ t ) [A2].
Given that Eq.
(15) corresponds to the product of Gaussian densities, the resulting normalized transition is also Gaussian, with mean and variance given by
[Equation]
where each µ jt is obtained using Eq.
Prior correction term The formulation above ignores the fact that the bridging densities defined in
Eq.
(3) involve the prior p ( θ ) .
We use the method proposed by Sohl-Dickstein et al.
[30] to correct for this, which involves adding the term σ 2 t (1 − n )( T − t ) T ∇ θ log p ( θ ) to the mean µ t from Eq.
(The derivation for this is similar to the one above, and also requires setting the resulting transition kernel to the normalized version of an unnormalized distribution [30].)
As mentioned previously, this derivation uses two assumptions/approximations.
[A1] assumes that the learned score function/reverse diffusion approximately reverses the noising process, which is reasonable if the forward kernels q t add small amounts of noise per step (equivalently, if the noise levels γ 1 , . .
, γ T increase slowly).
[A2] assumes that the normalized version of ˜ p t ( θ t − 1 | θ t ) , given by ˜ p
Nt ( θ t − 1 | θ t ) , approximately satisfies Eq.