Black Box Variational Inference (BBVI)

While Coordinate Ascent Variational Inference (CAVI) provides analytical updates for the variational parameters $\lambda$ when using conjugate models and the mean-field approximation, and Stochastic Variational Inference (SVI) scales this using stochastic gradients, both methods often rely on deriving model-specific, analytical expressions for the Evidence Lower Bound (ELBO) gradients. What happens when deriving these gradients analytically is difficult or even impossible? This situation arises frequently with:

• Non-conjugate models: When the prior and likelihood do not form a conjugate pair, the updates for the variational parameters in CAVI might not have a closed-form solution.
• Complex model structures: Models involving intricate dependencies, simulators, or functions within the generative process often make the joint probability $p(\mathbf{x}, \mathbf{z})$ difficult to differentiate analytically with respect to the variational parameters $\lambda$.

Black Box Variational Inference (BBVI) provides a general approach to VI that overcomes this limitation. It allows us to perform variational inference even when we cannot compute the gradient of the ELBO, $\nabla_{\lambda} \mathcal{L}(\lambda)$, analytically. The core idea is to approximate this gradient using Monte Carlo estimation. BBVI treats the model itself as a "black box," requiring only the ability to evaluate the log joint probability $\log p(\mathbf{x}, \mathbf{z})$ and the log variational density $\log q(\mathbf{z}|\lambda)$, along with their gradients under certain conditions.

Estimating the ELBO Gradient

Recall the ELBO: $\mathcal{L}(\lambda) = \mathbb{E}_{q(\mathbf{z}|\lambda)} [\log p(\mathbf{x}, \mathbf{z}) - \log q(\mathbf{z}|\lambda)]$

We need to compute $\nabla_{\lambda} \mathcal{L}(\lambda)$. The challenge lies in the fact that the expectation is taken with respect to $q(\mathbf{z}|\lambda)$, which itself depends on $\lambda$.

The Score Function Gradient Estimator

One way to proceed is using the log-derivative trick, also known as the score function identity: $\nabla_{\lambda} q(\mathbf{z}|\lambda) = q(\mathbf{z}|\lambda) \nabla_{\lambda} \log q(\mathbf{z}|\lambda)$. Applying this to the ELBO gradient yields:

$\nabla_{\lambda} \mathcal{L}(\lambda) = \nabla_{\lambda} \int q(\mathbf{z}|\lambda) [\log p(\mathbf{x}, \mathbf{z}) - \log q(\mathbf{z}|\lambda)] d\mathbf{z}$ $= \int \nabla_{\lambda} q(\mathbf{z}|\lambda) [\log p(\mathbf{x}, \mathbf{z}) - \log q(\mathbf{z}|\lambda)] d\mathbf{z} + \int q(\mathbf{z}|\lambda) \nabla_{\lambda} [\log p(\mathbf{x}, \mathbf{z}) - \log q(\mathbf{z}|\lambda)] d\mathbf{z}$

Assuming the model's joint probability $p(\mathbf{x}, \mathbf{z})$ does not depend on the variational parameters $\lambda$, the gradient of $\log p(\mathbf{x}, \mathbf{z})$ with respect to $\lambda$ is zero. The expression simplifies after applying the log-derivative trick and further manipulation:

$\nabla_{\lambda} \mathcal{L}(\lambda) = \mathbb{E}_{q(\mathbf{z}|\lambda)} [ (\nabla_{\lambda} \log q(\mathbf{z}|\lambda)) (\log p(\mathbf{x}, \mathbf{z}) - \log q(\mathbf{z}|\lambda)) ]$

This expectation can be approximated using Monte Carlo samples: draw $S$ samples $\mathbf{z}^{(s)} \sim q(\mathbf{z}|\lambda)$ and compute:

$\nabla_{\lambda} \mathcal{L}(\lambda) \approx \frac{1}{S} \sum_{s=1}^S (\nabla_{\lambda} \log q(\mathbf{z}^{(s)}|\lambda)) (\log p(\mathbf{x}, \mathbf{z}^{(s)}) - \log q(\mathbf{z}^{(s)}|\lambda))$

This is the score function gradient estimator, sometimes called the REINFORCE estimator in reinforcement learning contexts. To use it, we only need to be able to sample from $q(\mathbf{z}|\lambda)$, evaluate $\log p(\mathbf{x}, \mathbf{z})$, $\log q(\mathbf{z}|\lambda)$, and compute the gradient of $\log q(\mathbf{z}|\lambda)$ with respect to $\lambda$.

A significant drawback of the score function estimator is that it often exhibits high variance, requiring many samples $S$ for a stable estimate, which can slow down convergence.

The Reparameterization Trick (Pathwise Derivative Estimator)

An alternative, often lower-variance gradient estimator can be used when the variational distribution $q(\mathbf{z}|\lambda)$ is amenable to the reparameterization trick. This involves expressing the random variable $\mathbf{z}$ as a deterministic and differentiable transformation $g(\cdot, \lambda)$ of an auxiliary random variable $\epsilon$ with a fixed distribution $p(\epsilon)$ that does not depend on $\lambda$.

$\mathbf{z} = g(\epsilon, \lambda), \quad \text{where } \epsilon \sim p(\epsilon)$

For example, if $q(z|\lambda)$ is a Gaussian distribution $N(\mu, \sigma^2)$, where $\lambda = \{\mu, \sigma\}$, we can reparameterize $z$ as $z = \mu + \sigma \epsilon$, where $\epsilon \sim N(0, 1)$.

With this reparameterization, we can rewrite the ELBO expectation: $\mathcal{L}(\lambda) = \mathbb{E}_{p(\epsilon)} [\log p(\mathbf{x}, g(\epsilon, \lambda)) - \log q(g(\epsilon, \lambda)|\lambda)]$ Now, the expectation is with respect to $p(\epsilon)$, which is independent of $\lambda$. We can move the gradient inside the expectation: $\nabla_{\lambda} \mathcal{L}(\lambda) = \mathbb{E}_{p(\epsilon)} [\nabla_{\lambda} (\log p(\mathbf{x}, g(\epsilon, \lambda)) - \log q(g(\epsilon, \lambda)|\lambda))]$

The gradient can be approximated using Monte Carlo samples: draw $S$ samples $\epsilon^{(s)} \sim p(\epsilon)$, compute $\mathbf{z}^{(s)} = g(\epsilon^{(s)}, \lambda)$, and estimate: $\nabla_{\lambda} \mathcal{L}(\lambda) \approx \frac{1}{S} \sum_{s=1}^S \nabla_{\lambda} (\log p(\mathbf{x}, \mathbf{z}^{(s)}) - \log q(\mathbf{z}^{(s)}|\lambda))$ Using the chain rule, $\nabla_{\lambda} f(\mathbf{z}^{(s)})$ becomes $(\nabla_{\mathbf{z}} f(\mathbf{z}^{(s)}))^T (\nabla_{\lambda} g(\epsilon^{(s)}, \lambda))$. This requires gradients of the log joint and log variational density with respect to the latent variables $\mathbf{z}$, and the gradient of the transformation $g$ with respect to $\lambda$. Modern automatic differentiation libraries readily handle these computations.

The reparameterization trick typically yields gradient estimates with much lower variance compared to the score function method, often leading to faster and more stable convergence. However, it's not universally applicable; it requires finding a suitable function $g(\epsilon, \lambda)$.

Basic flow diagram illustrating the reparameterization trick within BBVI. Samples from a base distribution $\epsilon$ are transformed using variational parameters $\lambda$ to get samples $\mathbf{z}$. These are used to evaluate log probabilities, whose gradients with respect to $\lambda$ form the ELBO gradient estimate, which is then fed into an optimizer to update $\lambda$.

The BBVI Algorithm

BBVI leverages these gradient estimators within a stochastic optimization framework, similar to SVI.

1. Initialize variational parameters $\lambda$.
2. Repeat until convergence: a. Sample: Draw $S$ Monte Carlo samples. * If using the score function: Sample $\mathbf{z}^{(s)} \sim q(\mathbf{z}|\lambda)$ for $s=1,...,S$. * If using reparameterization: Sample $\epsilon^{(s)} \sim p(\epsilon)$ and compute $\mathbf{z}^{(s)} = g(\epsilon^{(s)}, \lambda)$ for $s=1,...,S$. b. Estimate Gradient: Compute the Monte Carlo estimate $\hat{\mathbf{g}}$ of $\nabla_{\lambda} \mathcal{L}(\lambda)$ using either the score function or reparameterization formula. If performing SVI, use the gradient estimate for the current data mini-batch. c. Update Parameters: Update $\lambda$ using a stochastic optimization algorithm (e.g., Adam, RMSprop) with the estimated gradient $\hat{\mathbf{g}}$: $\lambda_{t+1} \leftarrow \text{OptimizerUpdate}(\lambda_t, \hat{\mathbf{g}})$

When combined with data subsampling (as in SVI), step 2.b also incorporates stochasticity from the data mini-batch.

Advantages and Considerations

• Generality: BBVI's primary advantage is its applicability to a vast range of models, including those with non-conjugate structures or intractable analytic gradients. You only need to evaluate the model's log probabilities (and their gradients for reparameterization).
• Leverages Auto-Diff: Modern probabilistic programming languages (PPLs) like Pyro, NumPyro, TensorFlow Probability, and Stan heavily rely on automatic differentiation, making the implementation of BBVI (especially with the reparameterization trick) relatively straightforward for users. The PPL handles the gradient computations.
• Variance: The choice of gradient estimator is important. The score function estimator can suffer from high variance, potentially hindering convergence. The reparameterization trick, when applicable, is generally preferred for its lower variance. Advanced variance reduction techniques (e.g., control variates) can sometimes further improve the stability of both estimators.
• Tuning: Like other stochastic optimization methods, BBVI might require careful tuning of learning rates and optimizer parameters. The number of Monte Carlo samples $S$ per gradient step also influences the trade-off between gradient accuracy and computational cost.

BBVI significantly broadens the scope of models amenable to variational inference. By replacing analytical gradient derivations with Monte Carlo estimation, it provides a powerful and flexible tool for approximate Bayesian inference, particularly when combined with the scalability of stochastic optimization and the convenience of automatic differentiation.