Tutorial Part 3: Annealed Sequential-Monte-Carlo Sampler¶

Introduction¶

In this tutorial we using inference combinators to implement a Sequential Monte Carlo Sampler (Del Moral et al., 2006) that generates samples along a geometric annealing path

\begin{align} \gamma_k(x; \beta_k) := q_1(x)^{(1-\beta_k)}\gamma(x)^{\beta_k} ,&& \beta_1 = 0 ,&& \beta_k \in [0, 1] ,&& \beta_K = 1, \end{align}

as outlined in Zimmermann et al.. At each step $k$ we are given approximate samples $x_{k-1}$ from the last target $\gamma_{k-1}$ and use a variational proposals $q_k(\cdot \mid x_{k-1})$ to generate proposals $x_k$ for the next target density $\gamma_k$. We use these proposals to compute an importance weights

\begin{align} w_1 = 1 ,&& w_k = \frac{r_{k-1}(x_{k-1} \mid x_k)\gamma_k(x_k)}{\gamma_{k-1}(x_{k-1})q_k(x_k \mid x_{k-1})} && \text{for}\ 1 < k\leq K , \end{align}

which we can use to resample our particles in order to get approximate samples for the next target distribution $\gamma_k$.

While this sampling strategy is always valid, the quality of the sampler crucially depends on the variance of the importance weight above. The variance of the importance weights $w_k$ is minimized when

\begin{align} \check \gamma_k(x_{k-1}, x_k) := r_{k-1}(x_{k-1} \mid x_k)\gamma_k(x_k) && \text{and}\ && \hat \gamma_k(x_{k-1}, x_k) := \gamma_{k-1}(x_{k-1})q_k(x_k \mid x_{k-1}) \end{align}

are equal. To make $\hat \gamma_k$ as similar as possible to $\check\gamma_k$, we model the variational proposals $q_k$ and reverse kernels $r_k$ as variational distribution with parameters $\phi$ and $\theta$ respectively and minimize a KL-divergence

\begin{align} \mathcal{D}_{\mathrm{KL}}(\hat\gamma_{k} \mid \check \gamma_{k}) = \mathbb{E}_{\hat\gamma_{k}(x_{k-1}, x_{k}; \phi, \beta_{k-1})} \left[ \log \frac{\hat \gamma_k(x_{k-1}, x_k; \phi, \beta_{k-1})}{\check\gamma_k(x_{k-1}, x_k; \theta, \beta_{k})} \right] \end{align}

at each step. We additionally learn the parameters of the schedule of the annealing path $\beta_k$ (for $1<k<K$). Note that the optimization problem in this setting is underconstrained, however, the additional degrees of freedom might be helpful to make intermediate densities, between which the variational proposals struggle to map, close together and vice versa. For a more detailed exposition we refer to Zimmermann et al..

Defining the target and proposal density¶

We first define our target density which is an unnormalized gaussian mixture with $M=8$ modes, equally spaced around a circle with radius $10$ around the origin. Our initial proposal is Normal with zero mean and standard deviation $5$.

Defining a Sequence of annealed intermediate densities¶

Given the proposal and target densities defined above, we can define our intermediate annealed densities using flax. Note that we are defining the tunable parameters beta_raw in log space and normalize them appropriately to be in the interval $[0, 1]$.

Network¶

We now define our variational kernels, which each consist of a two-layer MLP with ReLU activations. We additionally define a helper class VariationalKernelList, which gives us convenient access to the individual variational kernels (automatically selecting the correct set of parameters), by providing the corresponding index. Lastly we define a Network class which wraps around the annealed target, forward kernels, and reverse kernels. This is convenient as it lets us pass around a single object which gives us access to the network of our individual model components.

Defining the model components¶

With all of our network definitions in place we can now define our model components as probabilistic programs in numpyro.

Using predefined inference algorithms in coix¶

Coix already implements a selection of inference algorithms including Nested Variational Inference (NVI). All we need to do is to instantiate our model components and pass it to the method that composes the inference program for us.

Evaluating the untrained model¶

As mentioned before, while our sampler might not be very efficient before we train it, it is still valid. So let's see how our sampler performs pre-training first.

While visually the samples loosely approximate the target density we can clearly see that the modes are too wide and not tightly enough peaked. This is also reflected in the statistics that we can access in the metics dictionary returned by the evaluation effect handler. We can see that the ess is around $300$ when taking $1000$ samples.

Training¶

So let's see if we can do better by training our model with NVI. All we need to do is to repeatedly run our inference program, differentiate the resulting loss, and take a gradient step. coix.util.train provides a convenient wrapper for such a training loop and optionally compiles the training procedure.

While the training loss with only 36 particles is not very informative to assess convergence, we can see that the ess goes up to about $34$ which is a good indicator that we learned a good proposal.

Evaluate trained model¶

We already implemented an evaluation wrapped, so let's evaluate the model again but this time with the optimized parameters.