Variational Autoencoders for Efficient Simulation-Based Inference

TL;DR

This work tackles likelihood-free simulation-based inference by proposing two variational autoencoder-based approaches: CP-VAE, which uses a data-dependent prior $p(\mathbf{z}\mid \mathbf{y})$ to adapt latent structure to observed data, and UP-VAE, which relies on an unconditional Gaussian prior $p(\mathbf{z})$ for simpler training. Both models aim to efficiently approximate complex posteriors $p(\boldsymbol{\theta} \mid \mathbf{y})$ in an amortized setting, demonstrated on sbibm benchmarks and Hodgkin–Huxley models. CP-VAE tends to offer greater flexibility through the conditional prior, while UP-VAE provides a stable, simpler alternative with competitive accuracy and lower risk of overfitting. The paper positions these VAEs as efficient, interpretable competitors to flow-based, GAN-based, and diffusion SBI methods, with future directions including integrating normalizing flows, richer priors, and architectural enhancements to capture more structure in the data.

Abstract

We present a generative modeling approach based on the variational inference framework for likelihood-free simulation-based inference. The method leverages latent variables within variational autoencoders to efficiently estimate complex posterior distributions arising from stochastic simulations. We explore two variations of this approach distinguished by their treatment of the prior distribution. The first model adapts the prior based on observed data using a multivariate prior network, enhancing generalization across various posterior queries. In contrast, the second model utilizes a standard Gaussian prior, offering simplicity while still effectively capturing complex posterior distributions. We demonstrate the ability of the proposed approach to approximate complex posteriors while maintaining computational efficiency on well-established benchmark problems.

Figures (8)

• Figure 1: CP-VAE Architecture. The encoder $q_{\psi}(\mathbf{z} \mid \mathbf{y}, \boldsymbol{\theta})$ computes the parameters of the variational distribution for latent variables $\mathbf{z}$, conditioned on paired simulated data $\mathbf{y}$ and parameters $\boldsymbol{\theta}$. The prior network $p(\mathbf{z} \mid \mathbf{y})$ defines a data-dependent prior over $\mathbf{z}$, enforcing latent space regularization via the KL divergence. The decoder $p(\boldsymbol{\theta} \mid \mathbf{y}, \mathbf{z})$ generates the distribution parameters for $\boldsymbol{\theta}$ based on $\mathbf{y}$ and $\mathbf{z}$. Both $\mathbf{z}$ and $\boldsymbol{\theta}$ are sampled using the reparameterization trick, ensuring end-to-end differentiability.
• Figure 2: UP-VAE architecture. The encoder $q_{\psi}(\mathbf{z} \mid \mathbf{y}, \boldsymbol{\theta})$ outputs the mean and variance of the approximate distribution, from which $\mathbf{z}$ is sampled (reparameterization trick). The KL divergence is computed between the encoder output and a standard normal prior. The theta decoder $p(\boldsymbol{\theta} \mid \mathbf{y}, \mathbf{z})$ outputs the parameters for $\boldsymbol{\theta}$, and the data decoder $p(\mathbf{y} \mid \mathbf{z})$ outputs the parameters for $\mathbf{y}$ (reparameterization trick).
• Figure 3: Two moons: true v/s estimated posterior (sim. budget: 30,000). (a) CP-VAE and (b) UP-VAE.
• Figure 4: SLCP : true v/s estimated posterior distribution (sim. budget: 30,000). (a) CP-VAE and (b) UP-VAE.
• Figure 5: HH Model: Estimated posterior (simulation budget: 10,000). (a) CP-VAE and (b) UP-VAE.