1. What variational Bayes is doing

Variational Bayes (VB) is a general framework for approximating a posterior distribution when:

  • The true posterior $p(\theta \mid y)$ is too hard to sample from directly (MCMC too slow or too big), or
  • Expectations like $E_p[h(\theta)]$ are too expensive to compute exactly.

Idea:

  • Choose a simpler family of distributions $g(\theta \mid \phi)$ indexed by variational parameters $\phi$.
  • Find the member of this family that is closest to the true posterior, in a specific sense (minimizing Kullback–Leibler divergence $\text{KL}(g \Vert p)$).

This is analogous to EM:

  • EM: alternates E-step and M-step to get a point estimate (mode / MLE / MAP).
  • VB: alternates updates of distributions to get a full approximate posterior.

2. Objective: minimize KL(g || p)

Target posterior: $p(\theta \mid y)$
Approximation: $g(\theta)$ (or $g(\theta \mid \phi)$)

The Kullback–Leibler divergence from $g$ to $p$ is:

$\text{KL}(g \Vert p) = – E_g\left[ \log \frac{p(\theta \mid y)}{g(\theta)} \right] = – \int \log \frac{p(\theta \mid y)}{g(\theta)} \, g(\theta) \, d\theta$

  • $\text{KL}(g \Vert p) \ge 0$.
  • $\text{KL}(g \Vert p) = 0$ only if $g(\theta) \equiv p(\theta \mid y)$.

We cannot usually optimize this exactly because:

  • We cannot evaluate or integrate $p(\theta \mid y)$ exactly in complex models.
  • But we often can evaluate the joint $p(\theta, y) = p(\theta)p(y \mid \theta)$ up to a constant.

So variational Bayes works with a form of this objective that uses $p(\theta, y)$ instead of $p(\theta \mid y)$.

3. Using a restricted class of approximating distributions

We choose a restricted class of approximations $g(\theta \mid \phi)$.

A common choice is mean-field factorization:

  • Let $\theta = (\theta_1, \dots, \theta_J)$.
  • Assume $g(\theta \mid \phi) = \prod_{j=1}^J g_j(\theta_j \mid \phi_j)$, i.e. the components $\theta_j$​ are independent under the approximation.

Then:

  • The form of each factor $g_j(\theta_j \mid \phi_j)$ is chosen based on the shape of the log posterior.
  • For exponential-family / conjugate models, each factor ends up being in a familiar family (Normal, Gamma, Inverse-$\chi^2$, etc.).

Key step:

  • For each component $\theta_j$​, consider:
    • $E_{g_{-j}}[\log p(\theta \mid y)]$ as a function of $\theta_j$​, where $g_{-j}$​ denotes the product of all other factors $\{g_k\}_{k \neq j}$​.
  • This expectation determines the functional form of $g_j$​.

You don’t compute the expectation numerically at this stage; you only identify the analytical form.

4. Variational Bayes update rule (mean-field)

Once the families $g_j(\theta_j \mid \phi_j)$ are identified, the algorithm proceeds as:

  1. Initialize all variational parameters $\phi = (\phi_1, \dots, \phi_J)$.
  2. Repeat until convergence:
    • For $j = 1, \dots, J$:
      • Update $g_j(\theta_j \mid \phi_j)$ by setting
        • $\log g_j(\theta_j \mid \phi_j) = E_{g_{-j}}[ \log p(\theta \mid y) ] + \text{constant}$
      • where the expectation is with respect to the current $g_{-j}(\theta_{-j})$.

Interpretation:

  • Each step updates one factor $g_j$​ assuming the others are fixed.
  • This is analogous to Gibbs sampling, but instead of sampling from the conditional, you set the whole distribution $g_j$​ to match the “best” exponential-family fit derived from expected log posterior.

For conjugate models, those expectations and resulting forms are closed-form.

5. Example: hierarchical 8-schools model

Model:

  • Parameters:
    • $\alpha_1, \dots, \alpha_8$​: school-specific effects (treatment effects)
    • $\mu$: population mean effect
    • $\tau$: population standard deviation of effects
  • Observed data: $y_j \sim N(\alpha_j, \sigma_j^2)$, with known $\sigma_j$.
  • Prior: $\alpha_j \sim N(\mu, \tau^2)$.

Let the full parameter vector be $\theta = (\alpha_1, \dots, \alpha_8, \mu, \tau)$

(10-dimensional).

The log posterior (up to constants) is:

$\log p(\theta \mid y) = -\frac{1}{2} \sum_{j=1}^8 \frac{(y_j – \alpha_j)^2}{\sigma_j^2} – 8 \log \tau – \frac{1}{2\tau^2} \sum_{j=1}^8 (\alpha_j – \mu)^2 + \text{const}$

Mean-field approximation:

$g(\theta) = g(\alpha_1) \cdots g(\alpha_8) g(\mu) g(\tau)$

5.1 Deriving the form of each factor

(a) Factor for each $\alpha_j$

Compute:

$E_{g_{-j}}[\log p(\theta \mid y)]$

as a function of $\alpha_j$​, averaging over:

  • all other $\alpha_k$​ for $k \neq j$,
  • $\mu$,
  • $\tau$.

Looking at the log posterior, the terms involving $\alpha_j$​ are:

$-\frac{1}{2} \frac{(y_j – \alpha_j)^2}{\sigma_j^2} – \frac{1}{2\tau^2} (\alpha_j – \mu)^2$

Taking expectation over $\mu$ and $\tau$ (under the current variational distributions):

  • The term remains quadratic in $\alpha_j$.
  • Therefore $g(\alpha_j)$ must be Normal.

The resulting form:

$g(\alpha_j) = N\left( \alpha_j \;\middle|\; \frac{y_j / \sigma_j^2 + E(1/\tau^2) E(\mu)} {1/\sigma_j^2 + E(1/\tau^2)}, \; \frac{1}{1/\sigma_j^2 + E(1/\tau^2)} \right)$

So the mean is a precision-weighted average of:

  • data information $y_j$ with precision $1/\sigma_j^2$
  • prior/pooled information $E(\mu)$ with “precision” $E(1/\tau^2)$.

(b) Factor for $\mu$

Consider expectations over $\alpha_1, \dots, \alpha_8, \tau$, keeping only terms that involve $\mu$:

$-\frac{1}{2\tau^2} \sum_{j=1}^8 (\alpha_j – \mu)^2$

Averaging over the current variational distributions gives a quadratic function in $\mu$:

$E[\log p(\theta \mid y)] \propto – \frac{1}{2} E\left(\frac{1}{\tau^2}\right) \sum_{j=1}^8 (E(\alpha_j) – \mu)^2 + \text{const}$

This is the log of a Normal density in $\mu$, so:

$g(\mu) = N\left( \mu \;\middle|\; \frac{1}{8} \sum_{j=1}^8 E(\alpha_j), \; \frac{1}{8 E(1/\tau^2)} \right)$

(c) Factor for $\tau^2$

Looking at the log posterior terms involving $\tau$:

$-8 \log \tau – \frac{1}{2\tau^2} \sum_{j=1}^8 (\alpha_j – \mu)^2$

This is equivalent to an inverse-gamma / scaled inverse-$\chi^2$ form. The variational factor is:

$g(\tau^2) = \text{Inv-}\chi^2\left( \tau^2 \;\middle|\; \nu = 7, \text{scale} = \frac{1}{7} \sum_{j=1}^8 E\big[(\alpha_j – \mu)^2\big] \right)$

So the approximate conditional distribution is:

$\tau^2 \sim \text{Inv-}\chi^2\left(7, M_\tau^2\right)$

where $M_\tau^2$​ is based on the expected squared deviations of $\alpha_j$​ from $\mu$ under the current variational distributions.

5.2 Rewriting in generic parameter form

We express the variational factors generically:

  • For each $j$, $g(\alpha_j) = N(\alpha_j \mid M_{\alpha_j}, S_{\alpha_j}^2)$
  • For $\mu$, $g(\mu) = N(\mu \mid M_\mu, S_\mu^2)$
  • For $\tau^2$, $g(\tau^2) = \text{Inv-}\chi^2(\tau^2 \mid 7, M_\tau^2).$

We then need to compute expectations like:

  • $E(\mu) = M_\mu$​,
  • $E(\alpha_j) = M_{\alpha_j}$​​,
  • $E(1/\tau^2) = 1/M_\tau^2$​,
  • $E[(\alpha_j – \mu)^2] = (M_{\alpha_j} – M_\mu)^2 + S_{\alpha_j}^2 + S_\mu^2$

These expectations are used to update the parameters in the next iteration.

5.3 Initialization and running the algorithm

Initialization

  • Choose starting values for the variational parameters, not for the true parameters.
    • For example:
      • $M_{\alpha_j} \sim N(0, 1)$,
      • Mμ∼N(0,1)M_\mu \sim N(0, 1)Mμ​∼N(0,1),
      • $S_{\alpha_j}, S_\mu \sim \text{Uniform}(0, 1)$.

Iteration

  1. For each $j = 1, \dots, 8$, update $g(\alpha_j)$ using the current $E(\mu)$ and $E(1/\tau^2)$.
  2. Update $g(\mu)$ using the current $E(\alpha_j)$ and $E(1/\tau^2)$.
  3. Update $g(\tau^2$ using current $E[(\alpha_j – \mu)^2]$.

Repeat until the parameters $M_{\alpha_j}, S_{\alpha_j}, M_\mu, S_\mu, M_\tau$ stabilize.

Plots show:

  • Each parameter’s path over iterations.
  • The KL divergence decreasing over iterations.

It often takes several tens of iterations to stabilize.

6. Monitoring convergence via the variational lower bound

Starting point:

$\text{KL}(g \Vert p) = -E_g\left[ \log \frac{p(\theta \mid y)}{g(\theta)} \right]$

Rewrite $p(\theta \mid y)$ in terms of the joint:

  • $p(\theta \mid y) = \frac{p(\theta, y)}{p(y)}$​.
  • So $\log p(\theta \mid y) = \log p(\theta, y) – \log p(y)$.

Then:

$\text{KL}(g \Vert p) = -E_g[\log p(\theta, y)] + E_g[\log g(\theta)] + \log p(y)$

Since $\log p(y)$ does not depend on $g$, minimizing $\text{KL}(g\Vert p)$ is equivalent to maximizing:

$\mathcal{L}(g) = E_g[\log p(\theta, y)] – E_g[\log g(\theta)]$

This is often called the variational lower bound or ELBO (Evidence Lower BOund).

In the 8-schools example, this is calculated explicitly (up to constants), and one can see that it increases (and thus KL decreases) each iteration.

7. Why each VB step decreases KL(g || p)

Consider updating just one factor $g_j(\theta_j)$ while keeping all others $g_{-j}$​ fixed.

Factorization:

$g(\theta) = g_j(\theta_j) g_{-j}(\theta_{-j})$

Rewrite the lower bound in a form focusing on $g_j$​. After algebra, the relevant part becomes:

$- \int g_j(\theta_j) \log \frac{\tilde{p}(\theta_j)}{g_j(\theta_j)} d\theta_j + \text{const}$

where $\log \tilde{p}(\theta_j) = E_{g_{-j}}[\log p(\theta, y)] + \text{const}.$

The quantity above is exactly KL$(g_j \Vert \tilde{p})$ up to a sign and an additive constant.

Thus:

  • Updating $g_j$​ to be proportional to $\tilde{p}$​ minimizes $\text{KL}(g_j \Vert \tilde{p})$.
  • This maximizes the lower bound $\mathcal{L}(g)$.
  • Therefore each update step decreases the global $\text{KL}(g \Vert p)$.

So as long as each update truly sets $\log g_j(\theta_j) = E_{g_{-j}}[\log p(\theta, y)] + \text{const},$

the algorithm is guaranteed to improve the approximation at every step.

8. Model checking under variational Bayes

Because $g(\theta)$ is a proper probability distribution:

  1. Draw $\theta^{(s)} \sim g(\theta)$ for $s = 1, \dots, S$.
  2. For each draw, simulate replicated data $y_{\text{rep}}^{(s)}$​ from the data model $p(y \mid \theta^{(s)})$.
  3. Compare $y_{\text{rep}}$​ to observed $y$.

This can detect:

  • obvious misfit in the generative model,
  • severe underestimation or overestimation of uncertainty.

Note: Since variational Bayes tends to underestimate variance, it can overfit as a point estimate of distributions, so model checking is still valuable.

9. Using VB as a starting point for better methods

VB is fast, but the approximation may be:

  • too narrow (underdispersed),
  • missing dependency structure.

To improve:

9.1 Importance sampling

  • Draw $\theta^{(s)} \sim g(\theta)$.
  • Compute weights $w^{(s)} \propto \frac{p(\theta^{(s)} \mid y)}{g(\theta^{(s)})}.$
  • Use weighted averages.

Problem: If $g$ is too narrow, importance weights can have enormous variance, making estimates unstable.

9.2 Importance resampling

  • Draw a large initial sample from $g$.
  • Compute weights as above.
  • Resample without replacement according to weights.
  • Use resampled points as approximate draws from $p(\theta \mid y)$.

This can be more stable than raw importance weighting but still delicate.

9.3 Particle filtering / MCMC refinement

  • Start from draws $\theta^{(s)}$ from $g$.
  • Use them as initial points for:
    • Metropolis,
    • HMC,
    • particle filters (with splitting/merging).
  • This can significantly reduce burn-in and improve efficiency.

10. EM as a special case of variational Bayes

EM also alternates between two objects, which we can reinterpret as variational factors:

  • Split parameters into $(\phi, \gamma)$.
  • Think of variational approximation:
    • $g(\phi)$: constrained to be a point mass (a Dirac delta at current estimate).
    • $g(\gamma)$: unconstrained; in fact, set equal to the exact conditional posterior:
      • $g(\gamma) = p(\gamma \mid \phi, y)$

Then:

  • Updating $g(\gamma)$ is just computing the conditional posterior given the current $\phi$ (E-step).
  • Updating $g(\phi)$ is equivalent to maximizing the lower bound with respect to $\phi$ (M-step).

So EM is a very special limiting case of variational inference where:

  • One factor is deterministic (point estimate),
  • The other factor is exact and unconstrained.

11. More general forms of variational Bayes

The version described so far:

  • Uses KL(g || p),
  • Uses a mean-field factorization,
  • Relies on conjugacy to get clean closed-form updates.

Extensions:

  • Approximate the objective (approximate KL) when conjugacy is not available.
    • Examples: logistic regression, generalized linear models.
  • Restrict $g(\theta)$ to e.g. multivariate Gaussians with full covariance, normalizing flows, etc.
  • Adapt the approximation at each step based on current $g$.

The core idea stays the same:

  • Choose a parametric family $g(\theta \mid \phi)$.
  • Iteratively update $\phi$ to minimize some approximation to $\text{KL}(g \Vert p)$.