STAT 238 - Bayesian Statistics Lecture Thirty Seven

Spring 2026, UC Berkeley

Variational Inference¶

We looked at the basics of variational inference in the last lecture. Consider the standard Bayesian setup with data given by $y$ and parameters given by $\theta$ . The prior is $f_{\theta}(\theta)$ and the likelihood is $f_{y \mid \theta}(y)$ . The posterior is then:

\begin{align*} f_{\theta \mid y}(\theta) = \frac{f_{\theta}(\theta) f_{y \mid \theta}(y)}{f_y(y)} \end{align*}

(1)

where the denominator is the evidence $f_{y}(y) = \int f_{y \mid \theta}(y) f_{\theta}(\theta) d\theta$ .

The most important quantity in variational inference is the ELBO defined as:

\begin{align} \text{ELBO}(q) := \int q(\theta) \log \frac{f_{y, \theta}(y, \theta)}{q(\theta)} d\theta \end{align}

(2)

where $q$ is an arbitrary probability density. The key facts about the ELBO are:

\begin{align} \text{ELBO}(q) \leq \log f_{y}(y) ~~ \text{ for every } q. \end{align}

(3)

and

\begin{align} \max_{q} \text{ELBO}(q) = \log f_y(y) ~~ \text{ with the maximum achieved at } q(\theta) = f_{\theta \mid y}(\theta). \end{align}

(4)

Both these facts are consequences of:

\begin{align*} \text{ELBO}(q) = \log f_{y}(y) - KL\left(q \| f_{\theta \mid y} \right). \end{align*}

(5)

Another useful fact about the ELBO (which follows directly from the definition (8) when $f_{y, \theta}(y, \theta)$ is replaced by $f_{y \mid \theta}(y) f_{\theta}(\theta)$ ) is:

\begin{align*} \text{ELBO}(q) = \int q(\theta) \log f_{y \mid \theta}(y) d\theta - KL(q \| f_{\theta}). \end{align*}

(6)

Connection to the EM Algorithm¶

The relation $\max_q \text{ELBO}(q) = \log f_{y}(y)$ (with the maximum being achieved at the posterior $q(\theta) = f_{\theta \mid y}(\theta)$ ) can be used to understand the EM algorithm.

Consider the same Bayesian setting as in the previous section, but suppose now that there is an additional hyperparameter $\alpha$ , and that both the likelihood $f_{y \mid \theta, \alpha}(y)$ as well as the prior $f_{\theta \mid \alpha}(\theta)$ depend on $\alpha$ .

The ELBO now becomes:

\begin{align*} \text{ELBO}(q, \alpha) = \int q(\theta) \log \frac{f_{y, \theta \mid \alpha}(y, \theta)}{q(\theta)} d\theta = \int q(\theta) \log \frac{f_{y \mid \theta, \alpha}(y) f_{\theta \mid \alpha}(\theta)}{q(\theta)} d\theta \end{align*}

(7)

and equation (4) now becomes:

\begin{align} \max_{q} \text{ELBO}(q, \alpha) = \log f_{y\mid \alpha}(y) ~~ \text{ with the maximum achieved at } q(\theta) = f_{\theta \mid y, \alpha}(\theta). \end{align}

(8)

Suppose now that we want to obtain the MLE $\hat{\alpha}$ of $\alpha$ i.e., we want to solve the optimization:

\begin{align*} \underset{\alpha}{\text{Maximize}}~ \log f_{y \mid \alpha}(y). \end{align*}

(9)

Using (8), we can rewrite the optimization above as:

\begin{align*} \underset{\alpha}{\text{Maximize}}~ \max_{q} \text{ELBO}(q, \alpha), \end{align*}

(10)

which can be viewed as the two parameter maximization problem:

\begin{align*} \underset{q, \alpha}{\text{Maximize}} ~ \text{ELBO}(q, \alpha). \end{align*}

(11)

It is natural to solve this two-parameter optimization alternatively. Fix $\alpha$ and optimize over $q$ , and then fix $q$ and optimize over $\alpha$ . This leads to the following algorithm.

Initialize $\alpha = \alpha^{(0)}$ .
Repeat the following for $t = 0, 1, 2, \dots$ (until some kind of convergence):
1. Take $q^{(t)}$ to be the maximizer of $\text{ELBO}(q, \alpha^{(t)})$ over all $q$ (for fixed $\alpha^{(t)}$ )
2. Take $\alpha^{(t+1)}$ to be the maximizer of $\text{ELBO}(q^{(t)}, \alpha)$ over all $\alpha$ .

By (8), it is clear that $q^{(t)}$ which maximizes $\text{ELBO}(q, \alpha^{(t)})$ is simply equal to:

q^{(t)} = f_{\theta \mid y, \alpha^{(t)}}.

Thus the alternating minimization algorithm above is equivalent to:

Initialize $\alpha = \alpha^{(0)}$ .
Repeat the following for $t = 0, 1, 2, \dots$ (until some kind of convergence):
1. Take $\alpha^{(t+1)}$ to be the maximizer of $\text{ELBO}(f_{\theta, \mid y, \alpha^{(t)}}, \alpha)$ over all $\alpha$ .

The quantity $\text{ELBO}(f_{\theta, \mid y, \alpha^{(t)}}, \alpha)$ equals:

\begin{align*} \text{ELBO}(f_{\theta, \mid y, \alpha^{(t)}}, \alpha) &= \int f_{\theta \mid y, \alpha^{(t)}}(\theta) \log \frac{f_{y, \theta \mid \alpha}(y, \theta)}{f_{\theta \mid y, \alpha^{(t)}}(\theta)} d\theta \\ &= \int f_{\theta \mid y, \alpha^{(t)}}(\theta) \log f_{y, \theta \mid \alpha}(\theta) d\theta - \int f_{\theta \mid y, \alpha^{(t)}}(\theta) \log f_{\theta \mid y, \alpha^{(t)}}(\theta) d\theta. \end{align*}

(12)

The second term actually does not depend on $\alpha$ , so maximizing $\text{ELBO}(f_{\theta, \mid y, \alpha^{(t)}}, \alpha)$ over $\alpha$ is equivalent to maximizing the first term above:

\begin{align} E(t, \alpha) := \int f_{\theta \mid y, \alpha^{(t)}}(\theta) \log f_{y, \theta \mid \alpha}(\theta) d\theta. \end{align}

(13)

The alternating maximization algorithm above then becomes:

Initialize $\alpha = \alpha^{(0)}$ .
Repeat the following for $t = 0, 1, 2, \dots$ (until some kind of convergence):
1. Calculate $E(t, \alpha)$ given in (13)
2. Take $\alpha^{(t+1)}$ to be the maximizer of $E(t, \alpha)$ over all $\alpha$ .

This is the EM algorithm. The calculation of $E(t, \alpha)$ is called the $E$ -step and the maximization of $E(t, \alpha)$ over $\alpha$ is called the $M$ -step. For more on this connection between Variational Inference (VI) and EM, see neal1998view.

Coordinatewise Variational Inference¶

The key problem that we need to solve in VI is that of maximizing $\text{ELBO}(q)$ . The full maximizer of $\text{ELBO}(q)$ over $q$ is, as already mentioned, $f_{\theta \mid y}$ . But this posterior is intractable in general and the very reason for using variational inerference is this intractability. The strategy is to restrict the class of $q$ over which we look for minimizers. This hopefully will lead to a tractable minimizer which is close to the actual posterior $f_{\theta \mid y}$ .

Suppose $\theta$ can be decomposed into $k$ separate classes of parameters as $\theta = (\theta_1, \dots, \theta_k)$ . We look for variational approximations of the form $q(\theta) = q_1(\theta_1) \dots q_k(\theta_k)$ . The ELBO maximization problem is:

\underset{q_1, \dots, q_k}{\text{maximize}}~ \text{ELBO}(q_1, \dots, q_k) = \int \dots \int q_1(\theta_1) \dots q_k(\theta_k) \log \frac{f_{y, \theta}(y, \theta)}{q_1(\theta_1) \dots q_k(\theta_k)} d\theta.

(14)

Alternating minimization leads to the following result.

The CAVI update equation (15) is equivalent to the following:

\begin{align} \log q_j^*(\theta_j) = \text{constant} + \E_{\theta_{-j}} \log f_{y, \theta_1, \dots, \theta_k}(y, \theta_1, \dots, \theta_k) \end{align}

(16)

where the expectation is taken with respect to $\theta_{-j} \sim q_{-j}$ .

Proposition 1 is a consequence of the following simple fact.

Lemma 1

Suppose $g$ is a nonnegative function that is not necessarily a density function. Then the minimizer of
$L(f) := \int f \log \frac{f}{g}$
over all densities $f$ is given by
$f^*(x) \propto g(x)$
(17)
which of course means that $f^* = g/(\int g)$ .
Suppose $f_2(\cdot)$ below is a fixed probability density function, and $g(x_1, x_2)$ is a fixed nonnegative function that is not necessarily a density. Then the minimizer of
$\Gamma(f_1) := \int \int f_1(x_1) f_2(x_2) \log \frac{f_1(x_1)f_2(x_2)}{g(x_1, x_2)} dx_1 dx_2$
over all densities $f_1$ is given by
$f_1^*(x_1) \propto \exp \left( \int f_2(x_2) \log g(x_1, x_2) dx_2 \right).$
(18)

Simple CAVI Example¶

As a simple illustration of how CAVI works, let us consider the problem of estimating $\mu, \sigma$ from observations $y_1, \dots, y_n$ that are i.i.d $N(\mu, \sigma^2)$ . As usual, we use the prior

\begin{align*} \mu, \log \sigma \overset{\text{i.i.d}}{\sim} \text{uniform}(-\infty, \infty). \end{align*}

(19)

In other words, the prior is $f_{\mu, \sigma^2}(\mu, \sigma^2)=\frac{1}{\sigma^2}$ (note that we are treating $\sigma^2$ and not $\sigma$ as the parameter).

The joint density of the data $y = (y_1, \dots, y_n)$ and the parameter $\theta = (\mu, \sigma^2)$ is given by:

\begin{align*} f_{y, \theta}(y, \theta) = \left(\prod_{i=1}^n \frac{1}{\sqrt{2 \pi} \sigma} \exp \left(-\frac{(y_i - \mu)^2}{2 \sigma^2} \right) \right) \frac{1}{\sigma^2} \end{align*}

(20)

so that

\begin{align*} \log f_{y, \theta}(y, \theta) &= -\sum_{i=1}^n \frac{(y_i - \mu)^2}{2 \sigma^2} - \left(\frac{n}{2} + 1 \right) \log \sigma^2 + \text{constant} \\ &= -\frac{n}{2 \sigma^2} (\bar{y} - \mu)^2 - \frac{1}{2 \sigma^2} \sum_{i=1}^n (y_i - \bar{y})^2 - \left(\frac{n}{2} + 1 \right) \log \sigma^2 + \text{constant}. \end{align*}

(21)

CAVI approximates the posterior of $\theta = (\mu, \sigma^2)$ by a product distribution $q_{\mu}(\mu) q_{\sigma^2}(\sigma^2)$ . Below we figure out how to update $q_{\mu}$ given $q_{\sigma^2}$ and also $q_{\sigma^2}$ given $q_{\mu}$ .

By (16), given $q_{\sigma^2}$ , we get

\begin{align*} \log q^*_{\mu}(\mu) = -\frac{n}{2} (\bar{y} - \mu)^2 \E_{q_{\sigma^2}} \left(\frac{1}{\sigma^2} \right) + \text{constant}, \end{align*}

(22)

which means that $q^*_{\mu}$ is the density of the normal distribution: with mean $\mu$ and variance

\begin{align} q^*_{\mu} = \text{ density of } N\left(\bar{y}, \frac{1}{n \E_{q_{\sigma^2}}(1/\sigma^2)} \right). \end{align}

(23)

To get $q^*_{\sigma^2}(\sigma^2)$ for fixed $q_{\mu}$ , we again use (16) to get

\begin{align*} \log q^*_{\sigma^2}(\sigma^2) = -\frac{1}{2 \sigma^2} \left(n \E_{q_{\mu}} (\bar{y} - \mu)^2 + \sum_{i=1}^n (y_i - \bar{y})^2 \right) - \left(\frac{n}{2} + 1 \right) \log \sigma^2 + \text{constant}. \end{align*}

(24)

This has the form of the Inverse Gamma density $IG(a, b)$ given by:

\begin{align*} -\frac{b}{\sigma^2} - (a + 1) \log \sigma^2. \end{align*}

(25)

Comparing terms, we get

\begin{align} q^*_{\sigma^2} = \text{ density of } IG \left(\frac{n}{2}, \frac{n}{2} \E_{q_{\mu}} (\bar{y} - \mu)^2 + \frac{1}{2} \sum_{i=1}^n (y_i - \bar{y})^2 \right). \end{align}

(26)

(23) and (26) give the following CAVI algorithm:

Initialize with some density $q^{(0)}_{\sigma^2}$ of $\sigma^2$ . Suppose $s_0 := \left(n \E_{q^{(0)}_{\sigma^2}}(1/\sigma^2) \right)^{-1}$ .
Repeat the following for $k = 0, 1, 2, \dots$ :
1. Take $q_{\mu}^{(k+1)}$ to be the density of $N(\bar{y}, s_k^2)$ .
2. Take $q_{\sigma^2}^{(k+1)}$ to be the density of
  $\begin{align*} IG \left(\frac{n}{2}, \frac{n}{2}\E_{q^{(k+1)}_{\mu}} (\bar{y} - \mu)^2 + \frac{1}{2} \sum_{i=1}^n (y_i - \bar{y})^2 \right) = IG \left(\frac{n}{2}, \frac{n}{2}s_k^2 + \frac{1}{2} \sum_{i=1}^n (y_i - \bar{y})^2 \right) \end{align*}$
  (27)
3. Take $s_{k+1}^2 = \left(n \E_{q_{\sigma^2}^{(k+1)}} (1/\sigma^2) \right)^{-1}$ which leads to (using the fact $\E \sigma^{-2} = a/b$ when $\sigma^2 \sim IG(a, b)$ )
  $\begin{align} s_{k+1}^2 = \frac{s_k^2}{n} + \frac{1}{n^2} \sum_{i=1}^n (y_i - \bar{y})^2. \end{align}$
  (28)

The limit of the iteration for $\{s_k^2\}$ in (28) can be computed in closed form. Indeed if $s^2 = \lim_{k \rightarrow \infty} s_k^2$ , then (28) gives

\begin{align*} s^2 = \frac{s^2}{n} + \frac{1}{n^2} \sum_{i=1}^n (y_i - \bar{y})^2 \implies s^2 = \frac{1}{n(n-1)} \sum_{i=1}^n (y_i - \bar{y})^2. \end{align*}

(29)

Thus the VI solution $q^*_{\mu}(\mu) q^*_{\sigma^2}(\sigma^2)$ is given by:

\begin{align} q_{\mu}^* = \text{ density of } N \left(\bar{y}, \frac{1}{n(n-1)} \sum_{i=1}^n (y_i - \bar{y})^2 \right) \end{align}

(30)

and

\begin{align} q_{\sigma^2}^* &= \text{ density of } IG\left(\frac{n}{2}, \frac{1}{2(n-1)} \sum_{i=1}^n (y_i - \bar{y})^2 + \frac{1}{2} \sum_{i=1}^n (y_i - \bar{y})^2 \right) \nonumber \\ &= \text{ density of } IG\left(\frac{n}{2}, \frac{n}{2(n-1)} \sum_{i=1}^n (y_i - \bar{y})^2 \right) \end{align}

(31)

We can compare these VI marginals to those of the exact posterior which we know is given by:

\begin{align*} \mu \mid \text{data} \sim t_{n-1} \left(\bar{y}, \frac{1}{n(n-1)} \sum_{i=1}^n (y_i - \bar{y})^2 \right) ~~ \text{ and } ~~ \sigma^2 \mid \text{data} \sim IG \left(\frac{n-1}{2}, \frac{1}{2} \sum_{i=1}^n (y_i - \bar{y})^2 \right). \end{align*}

(32)

Comparing the above to (30), we see that the VI marginal posterior has a narrower normal density instead of $t_{n-1}$ . Comparing the above to (31), we see that the parameters of the IG distribution are changed slightly (the changes are negligible for large $n$ ). Note also that $\mu$ and $\sigma^2$ are independent in the actual posterior but this VI analysis assumes they are independent.