Variational Inference

Kullback-Leibler Divergence is a non-symmetric measure of how one probability distribution diverges from a second, reference distribution. In VI, the reverse KL divergence, KL(q || p), is minimized.

• Reverse KL (q || p): KL(q(z) || p(z|x)) = ∫ q(z) log(q(z)/p(z|x)) dz. Minimizing this encourages q to be mode-seeking; it tends to concentrate on a single mode of p, potentially underestimating variance.
• Forward KL (p || q): The alternative would be mode-covering but is intractable as it requires expectations under p.
• Role as Regularizer: In the ELBO decomposition, the KL term acts as a regularizer, penalizing complex q distributions that deviate too far from the prior p(z).

The Reparameterization Trick is a key technique for enabling efficient, low-variance gradient estimation of the ELBO with respect to the variational parameters φ. It is essential for the Stochastic Gradient Variational Bayes (SGVB) algorithm.

• Mechanism: It expresses the random variable z ~ q(z; φ) as a deterministic function z = g(φ, ε) of the parameters φ and an auxiliary noise variable ε drawn from a fixed distribution (e.g., ε ~ N(0, I)).
• Gradient Estimation: This allows gradients of the ELBO to be estimated as ∇_φ 𝔼_q[f(z)] = 𝔼_p(ε)[∇_φ f(g(φ, ε))], which can be approximated with Monte Carlo samples. This yields gradients with much lower variance than the alternative score function estimator.
• Applicability: Used for continuous latent variables with distributions like the Gaussian. For discrete variables, alternative methods like the Gumbel-Softmax trick are employed.

Amortized Variational Inference scales VI to large datasets by learning a shared, parametric function (an inference network) that maps observations x directly to the parameters of the variational distribution q(z | x; φ).

• Inference Network: Typically a neural network that outputs the mean and variance for a Gaussian q. This is the core of the Variational Autoencoder (VAE).
• Efficiency: Instead of optimizing separate variational parameters φ_i for each data point x_i, a single network is trained. After training, approximate posterior inference for a new x is a single forward pass.
• Trade-off: The amortization gap is the discrepancy between the optimal per-datapoint variational parameters and the output of the inference network. A powerful network minimizes this gap.

Stochastic Optimization is the practical engine for maximizing the ELBO on large datasets. It uses mini-batches of data to compute noisy, unbiased gradient estimates, enabling scalable learning.

• Stochastic Gradient Variational Bayes (SGVB): The standard algorithm. For each mini-batch, it:
  1. Samples data points x.
  2. Uses the inference network (if amortized) to get q(z|x) parameters.
  3. Samples z via the reparameterization trick.
  4. Computes a Monte Carlo estimate of the ELBO gradient for the mini-batch.
  5. Updates all parameters (of both the generative model and inference network) using gradient ascent.
• Adaptive Optimizers: Algorithms like Adam are universally used due to their efficiency in tuning learning rates for the complex, high-dimensional optimization landscape of deep VI models.

A latent variable is an unobserved random variable that is inferred to explain the observed data. In models using variational inference, these variables (e.g., z) represent a compressed, explanatory code for the input (e.g., x). The latent space is the continuous, lower-dimensional manifold where these latent representations reside. Operations in this space are powerful:

• Interpolation: Smoothly moving between points can generate semantically meaningful transitions in data space.
• Disentanglement: Ideally, individual dimensions control independent factors of variation (e.g., object size, rotation, color).
• Sampling: New data is generated by sampling a latent vector z ~ p(z) and passing it through a generative decoder.

A generative model learns the joint probability distribution p(x, z) or the marginal p(x) of the observed data. Variational inference is a primary technique for training such models, especially when the posterior is intractable. Prominent examples include:

• Variational Autoencoder (VAE): A deep generative model that uses an encoder (inference network) to parameterize q(z|x) and a decoder to parameterize p(x|z), trained by maximizing the ELBO.
• Normalizing Flows: Uses a series of invertible transformations to map a simple distribution to a complex one, allowing for exact likelihood computation but often with more restrictive architectures. These models contrast with discriminative models, which learn the conditional distribution p(y|x).

A Bayesian Neural Network treats its weights as probability distributions rather than deterministic point estimates. This provides a principled framework for uncertainty quantification. Variational inference is the standard scalable method for approximating the posterior over these millions of weights.

• Epistemic Uncertainty: Captured by the distribution over weights; reflects model uncertainty due to limited data. It can be reduced with more data.
• Practical Training: The Bayes by Backprop algorithm uses variational inference, often with a mean-field Gaussian posterior, to learn weight distributions. Predictions are made by marginalizing over the weights, typically approximated via Monte Carlo sampling.

Amortized variational inference scales inference to large datasets by learning a shared, parameterized inference network (e.g., an encoder) that maps any input x directly to the parameters of its variational posterior q(z|x; φ). This contrasts with traditional VI, which optimizes a separate q for each individual data point.

• Efficiency: After training, inferring q(z|x) for new data is a single forward pass.
• Core to VAEs: The encoder network is the amortization vehicle.
• Potential Pitfall: Amortization gap refers to the sub-optimality introduced by using a shared function approximator instead of per-datum optimization, which can lead to a looser ELBO.