STAT 238 - Bayesian Statistics Lecture Thirty Eight

Spring 2026, UC Berkeley

As an example of variational inference, we shall look at an example of Variational Auto-Encoders (VAE) which are used for generative modeling.

Variational AutoEncoders¶

Consider data points $y_1, \dots, y_n$, where each $y_i$ is a binary vector in $\{0, 1\}^{d_y}$. One can think of each $y_i$ as an image, flattened from its 2D pixel grid into a single vector of length $d_y$, where each entry corresponds to one pixel. In general, pixel values can take many forms—for instance, grayscale images use intensities in $\{0, 1, \dots, 255\}$, and color images use three such values per pixel (one per RGB channel). Here we focus on the simplest case: binary images (also called black-and-white or bilevel images), in which each pixel is either off (0) or on (1). For example, a $28 \times 28$ binary image yields $d_y = 28^2 = 784$.

We build up a model for the data in stages. Since each coordinate of $y_i$ is binary, the most basic assumption is

where $p_i \in [0, 1]^{d_y}$ specifies the probability that each pixel is on (Bernoulli is interpreted componentwise, so pixels are conditionally independent given $p_i$). In other words, we are assuming that each $y_i \in \{0, 1\}^{d_y}$ is Bernoulli with its own parameter $p_i \in [0, 1]^{d_y}$. We need to further model $p_1, \dots, p_n$ (otherwise, there will be too many parameters to yield anything useful). Here it is more convenient to work on the logit scale: let

applied componentwise.

We shall assume $\eta_1, \dots, \eta_n$ are i.i.d with some common distribution on $\R^{d_y}$. Further, we assume that this common distribution, despite living in $\mathbb{R}^{d_y}$, is in fact concentrated on (or near) a much lower-dimensional set. A natural way to encode this is to write

where $f_\theta : \mathbb{R}^d \to \mathbb{R}^{d_y}$ is a neural network with parameters $\theta$ and $d \ll d_y$. The pushforward of $N(0, I_d)$ under $f_\theta$ defines a $d$-dimensional family of distributions on $\mathbb{R}^{d_y}$, and the flexibility of $f_\theta$ ensures this family is rich enough to model the structure we care about. The Gaussian prior on $z_i$ is essentially a convention: any continuous distribution on $\mathbb{R}^d$ can be written as a deterministic transformation of a standard Gaussian, so fixing $z_i \sim N(0, I_d)$ and letting $f_\theta$ be flexible loses no generality (it is also analytically convenient, as we will see when designing an inference procedure). Putting the pieces together, and writing $\sigma$ for the componentwise sigmoid, the overall model can be written in the following alternative way:

with, for example,

The unknown parameters in this model are the neural network parameters $\theta$. The latent variables $z_1, \dots, z_n$ are also unobserved. The $z_i$’s can be integrated out and the model becomes

where

with $\varphi(z)$ being the density of $N(0, I_d)$, and

The likelihood then is given by:

and the corresponding log-likelihood is

This log-likelihood cannot be directly used as the objective function in optimization software such as PyTorch because of the presence of the integral. A natural idea is to discretize the integral by sampling $z^{(l)}_i, l = 1, \dots, M$ from $N(0, I_d)$ and forming the approximate log-likelihood:

The above term (11) is unlikely to be a good approximation to the actual log-likelihood (10). This is because the integral $\int p_{\theta}(y_i \mid z_i) \varphi(z_i) dz_i$ in (10) is dominated by the region where $p_{\theta}(y_i \mid z_i)$ is large i.e., by the posterior $p_{\theta}(z_i \mid y_i)$. This posterior typically is very concentrated in the sense that only a tiny region of $z_i \in \R^d$ yields non-negligibe values of $p_{\theta}(y_i \mid z_i)$. The prior $z_i \sim N(0, I_d)$, by contrast, spreads mass diffusely over $\R^d$. The consequence will be that almost every sample $z_i^{(l)}$ lands in a region where $p_{\theta}(y_i \mid z_i^{(l)})$ is small, so that the average in (11) is dominated by one or two (or even zero) lucky samples. This makes the average in (11) a highly variance estimate of the integral in (10), and hence unreliable as an optimization objective.

A better approach is to use ideas from variational inference and to maximize the ELBO instead.

Use of the ELBO for maximization of $\log f_{\text{data} \mid \theta}(y_1, \dots, y_n)$ ¶

We shall use the following formula whose proof we saw in Lecture 36 (see also Lecture 37):

where

where $q_i$ is the $i$-th marginal corresponding to $q$, and $\varphi \times \dots \times \varphi$ denotes the distribution of $(z_1, \dots, z_n)$ corresponding to $z_i \overset{\text{i.i.d}}{\sim} N(0, I_d)$.

The maximization in (12) is over all probability densities $q$ of $z_1, \dots, z_n$. We have also seen (again in Lecture 36) that the maximum in (12) is achieved when $q$ is the conditional density of $z_1, \dots, z_n$ given $y_1, \dots, y_n$ and $\theta$:

It is clear that this $q^*$ is a product measure $\prod_{i=1}^n q_i^*(z_i \mid y_i)$ where

However because of the somewhat complicated form of $p_{\theta}(y_i \mid z_i)$, the density $q_i^*$ above is not of a simple form (for example, it is not a Gaussian). We therefore take a simpler class of densities $\Qcal$ and then maximize the ELBO over $\Qcal$. Specifically, we shall take $\Qcal$ to be the class of all product densities $\prod q_{i, \phi}(z_i \mid y_i)$ for which each marginal $q_{\phi}(z_i \mid y_i)$ is Gaussian:

with

With this choice $\Qcal$, the ELBO becomes

where we use the fact that the KL between two product distributions becomes the sum of the KL between the marginals. We can further write this as:

The second term above can be written in closed form as (see e.g., {https://en.wikipedia.org/wiki/Kullback

For the first term in the ELBO expression above, we need to compute

For this, we can use Monte Carlo averaging. First write $z_i \sim N(\mu_{\phi}(y_i), \Sigma_{\phi}(y_i))$ as

Here $\odot$ stands for the Hadamard product (elementwise multiplication). Note that the components of $\Sigma_{\phi}(y_i)^{1/2} u_i$ are simply $\sigma_{\phi, j}(y_i) u_{i, j}$. Thus

where $u_i^{(l)}, 1 \leq l \leq L, 1 \leq i \leq n$ are i.i.d samples generated from $N(0, I_d)$. Thus the ELBO becomes:

Here $\sigma$ denotes the sigmoid function while $\sigma_{\phi}$ denotes the variance parameter of the variational distribution $q_{\phi}$.

We simply maximize $\text{ELBO}(\phi, \theta)$ jointly over $\phi$ and $\theta$ to estimate them. We will parametrize $\mu_{\phi}$ and $\sigma_{\phi}$ using neural networks. For example,

with $\phi = (W_{\mu}, W, w, b_{\mu}, W_{\sigma}, b_{\sigma})$. Recall also that $f_{\theta}$ will also be parametrized by a neural network as in (5).

With these parametrizations, $\text{ELBO}(\phi, \theta)$ can be maximized using optimization software such as PyTorch.