Variational_inference.qmd

---
title: Supervised learning methods - Generative models
subtitle: Variational Inference and VAE's
format: clean-revealjs
title-slide-attributes:
    data-background-image: figs/logo_without_leaf-1.png
    data-background-size: 30%
    data-background-position: bottom 70px right 100px
author:
  - name: Xenofon Karakonstantis
    affiliations:
      - name: Demant
  - name: Samuel A. Verburg
    affiliations:
      - name: DTU Electro
date: today
bibliography: VIref.bib
link-citations: true
---

## Outline

- The problem of intractable posteriors  
- From inference to optimisation  
- The Evidence Lower Bound (ELBO)  
- Variational Bayes and VB-EM  
- Connection to Variational Autoencoders (VAEs)

---

## The Bayesian inference problem

We often want:
$$
p(\mathbf z|\mathbf x) = \frac{p(\mathbf x,\mathbf z)}{p(\mathbf x)} = \frac{p(\mathbf x|\mathbf z)p(\mathbf z)}{\int p(\mathbf x|\mathbf z)p(\mathbf z)\,d\mathbf z}
$$

But $p(\mathbf x)$ (the *evidence*) is often **intractable**:

- High-dimensional latent space  

- Non-conjugate priors or nonlinear likelihoods (not analytically tractable)

$$
\Rightarrow \text{Need approximate inference.}
$$

---

## Example: Bayesian Gaussian Mixture

Generative model
$$
\begin{aligned} \mu_k &\sim \mathcal N(0, \tau^2), \\ z_n &\sim \text{Multinomial}( \pi), \\ x_n | z_n &\sim \mathcal N( \mu_{ z_n}, \sigma^2).
\end{aligned}
$$

Goal: posterior $p(\pmb \mu, \mathbf{z} \mid \mathbf{x})=\frac{p(\pmb{\mu}, \mathbf{z}, \mathbf{x})}{\int_\mu \sum_{\mathbf{z}} p(\pmb \mu, \mathbf{z}, \mathbf{x})}$

Problem (marginal likelihood - denominator)
$$
p(\mathbf x) = \int_\mu \sum_z p(\pmb \mu, \mathbf z, \mathbf x) \quad \text{is intractable.}
$$

---

## Three families of approximate inference {.smaller}

1. **Sampling methods (MCMC):**
   - Asymptotically exact
   - Slow, difficult to diagnose convergence

2. **Variational inference (VI):**
   - Optimise a tractable distribution $q(\mathbf z|\pmb \nu)$
   - Fast, scalable to large datasets  

3. **Normalising Flows**^[@rezende2015variational]
   - Flexible posteriors via invertible transforms; exact log-density by change of variables
   - Can be used as an expressive $q(\mathbf z)$ inside VI (flow-based variational families)
   - Trade-off: higher compute; requires Jacobian-tractable layers (e.g., planar, RealNVP, Glow)

Maybe we can... [turn inference into optimisation task]{.fg style="--col: #3255a8ff"}

---

## Approximate inference {.smaller}

### Variational methods 
- Consider a family of tractable distributions $q(\mathbf z|\pmb \nu)$ - the variational distribution

- Find the variational parameters $\pmb \nu$ that make $q(\mathbf z|\pmb \nu)$ as close as possible to the true posterior $p(\mathbf z| \mathbf x)$

- Turns the inference problem into an optimisation problem!

- Use $q(\mathbf z|\pmb \nu)$ with the fitted parameters as a proxy for the true posterior
    - To make predictions about future data
    - To investigate the posterior distribution of the hidden variables

```{python}
#| label: fig-bimodal-toy
#| fig-cap: "Bimodal target and a simple Gaussian approximation (moment-matched)."
#| echo: false
#| warning: false

import numpy as np
import matplotlib.pyplot as plt

# Target mixture p(x) = 0.5 N(-2, 0.5^2) + 0.5 N(2, 0.5^2)
w = np.array([0.5, 0.1])
mus = np.array([-2.0,  0.5])
sigmas = np.array([0.5, 0.5])

def norm_pdf(x, mu, sigma):
    return np.exp(-0.5 * ((x - mu) / sigma)**2) / (np.sqrt(2*np.pi) * sigma)

def p_density(x):
    total = np.zeros_like(x, dtype=float)
    for wi, mui, si in zip(w, mus, sigmas):
        total += wi * norm_pdf(x, mui, si)
    return total

# Moment-matched Gaussian q(x) = N(mean, var)
mix_mean = -2.0
mix_var  = np.sum(w * (sigmas**2 + (mus - mix_mean)**2))
q_mu, q_sigma = mix_mean , np.sqrt(mix_var)

def q_density(x):
    return norm_pdf(x, q_mu, q_sigma)

# Plot densities
x = np.linspace(-6, 6, 1000)
px = p_density(x)
qx = q_density(x)

plt.figure(figsize=(7, 3.5))
plt.plot(x, px, label='Target p(x): bimodal', color='#1f77b4', lw=2)
plt.plot(x, qx, label='Approx q(x): Gaussian', color='#ff7f0e', lw=2, ls='--')
plt.fill_between(x, 0, px, color='#1f77b4', alpha=0.08)
plt.fill_between(x, 0, qx, color='#ff7f0e', alpha=0.08)
plt.xlabel('x'); plt.ylabel('density')
plt.legend(loc='upper right', frameon=False)
plt.tight_layout()
```
---

## Mean-field approximation

- How should we choose a tractable family of distributions for $q(\mathbf z)$?

    - If we assume independence between latent variables
    $$
    q(\mathbf z) = \prod_{m=1}^M q(\mathbf z_m),
    $$

this simplifies expectations and leads to **coordinate ascent** updates.

---

## Variational Inference {.smaller}

### Kullback-Leibler divergence

:::: {.columns}

::: {.column width="40%"}

Find variational parameters $\pmb \nu$ that make $q(\mathbf z|\pmb \nu)$ as close as possible to true posterior $p(\mathbf z| \mathbf x)$

$$
\pmb \nu^* = \arg\min_{\pmb \nu} \text{KL}(q(\mathbf z|\pmb \nu)\|p(\mathbf z|\mathbf x))
$$

where the **Kullback–Leibler divergence** is

$$
\begin{align}
\text{KL}(q\|p) &= \int_{\mathbf{z}} q(\mathbf{z}) \log \frac{q(\mathbf{z})}{p(\mathbf{z} \mid \mathbf{x})} \\ 
&= \mathbb{E}_q\left[\log \frac{q(\mathbf{z})}{p(\mathbf{z} \mid \mathbf{x})}\right]
\end{align}
$$

:::

::: {.column width="60%"}
![](figs/Blei_KLdiv.png){fig-cap="KL divergence geometry (adapted from Blei)" fig-alt="KL divergence between q and p visualised" .absolute bottom=150 right=10 width="550" height="300"}

:::

::::

::: aside
Figure from @blei2017variational
:::

## From KL minimisation to ELBO

Start with:
$$
\text{KL}(q(\mathbf z)\|p(\mathbf z|\mathbf x)) = \mathbb E_q[\log q(\mathbf z)] - \mathbb E_q[\log p(\mathbf z|\mathbf x)]
$$

Expand $p(\mathbf z|\mathbf x) = \frac{p(\mathbf z,\mathbf x)}{p(\mathbf x)}$:
$$
\text{KL}(q\|p) = \mathbb E_q[\log q(\mathbf z)] - \mathbb E_q[\log p(\mathbf z,\mathbf x)] + \log p(\mathbf x)
$$

Rearrange:
$$
\log p(\mathbf x) = \text{KL}(q\|p) + \underbrace{(\mathbb E_q[\log p(\mathbf z,\mathbf x)] - \mathbb E_q[\log q(\mathbf z)])}_{\mathcal{L}(q)}
$$

---

## The Evidence Lower Bound (ELBO)^[@jordan1999introduction @bishop2006pattern]

Define:
$$
\mathcal{L}(q) = \mathbb E_q[\log p(\mathbf z,\mathbf x)] - \mathbb E_q[\log q(\mathbf z)]
$$

Since KL ≥ 0:
$$
\mathcal{L}(q) \le \log p(\mathbf x)
$$

Maximising the ELBO is equivalent to minimising KL.

---

## Derivation via Jensen’s inequality

$$
\begin{aligned}
\log p(\mathbf x)
  &= \log \int p(\mathbf x,\mathbf z)\,d\mathbf z
   = \log \int q(\mathbf z) \frac{p(\mathbf x,\mathbf z)}{q(\mathbf z)}\,d\mathbf z \\
  &= \log \mathbb E_q\left[\frac{p(\mathbf x,\mathbf z)}{q(\mathbf z)}\right]
   \ge \mathbb E_q[\log p(\mathbf x,\mathbf z) - \log q(\mathbf z)] = \mathcal L(q)
\end{aligned}
$$

Equality holds when $q(\mathbf z) = p(\mathbf z|\mathbf x)$.

---

## Two equivalent views of the ELBO

1. **Expected log joint – entropy:**
   $$
    \mathcal{L}(q)=\mathbb{E}_q[\log p(\mathbf{z}, \mathbf{x})]-\underbrace{\mathbb{E}_q[\log q(\mathbf{z} \mid \boldsymbol{\nu})]}_{\text {entropy terms }}
   $$

2. **Reconstruction – regularisation:**
   $$
   \mathcal L(q) = \mathbb E_q[\log p(\mathbf x|\mathbf z)] - \text{KL}(q(\mathbf z)\|p(\mathbf z))
   $$

    - $\mathbb{E}_q[\log p(\mathbf{x} \mid \mathbf{z})]$ is the reconstruction loss

    - $\text{KL}(q(\mathbf{z}) \| p(\mathbf{z}))$ acts like a regularisation term (to stay close to the prior)

The latter is the form used in **VAEs**.

## Variational Bayes
### Connection to Variational Autoencoders (VAEs)

The VAE objective *is* the ELBO^[@kingma2013auto]
$$
\mathcal L(\theta,\phi) =
\mathbb E_{q_\phi(\mathbf z|\mathbf x)}[\log p_\theta(\mathbf x|\mathbf z)]
 - \text{KL}(q_\phi(\mathbf z|\mathbf x)\|p(\mathbf z))
$$

where  
- Encoder $q_\phi(\mathbf z|\mathbf x)$: inference network  
- Decoder $p_\theta(\mathbf x|\mathbf z)$: generative model

---

## Variational Autoencoders {.smaller}

### The reparameterisation trick

- Encoder (inference network) outputs $\pmb \mu_\phi(\mathbf x)$ and $\log \pmb \sigma^2_\phi(\mathbf x)$, parameterising
  $$
  q_\phi(\mathbf z|\mathbf x)=\mathcal N\!\big(\pmb \mu_\phi(\mathbf x), \operatorname{diag}(\pmb \sigma_\phi^2(\mathbf x))\big),
  $$
  
- $\mathbf z$ is the latent code sampled via reparameterisation
  $$
  \mathbf z=\pmb \mu_\phi(\mathbf x)+\pmb \sigma_\phi(\mathbf x)\odot \pmb \epsilon,\ \ \pmb \epsilon\sim\mathcal N(\mathbf 0,\mathbf I)
  $$.

This allows gradients to flow through $\mathbf z$.

$$
\nabla_\phi \mathcal L(\theta,\phi)
  = \nabla_\phi \mathbb E_{\epsilon}[\log p_\theta(\mathbf x|\mathbf z) - \text{KL}(q_\phi(\mathbf z|\mathbf x)\|p(\mathbf z))]
$$

## Variational Autoencoders

### Decoder and latent regularisation

- Decoder (generative network) takes $\mathbf z$ and outputs parameters of $p_\theta(\mathbf x|\mathbf z)$ (reconstruction).

- Prior $p(\mathbf z)=\mathcal N(\mathbf 0,\mathbf I)$; the KL term keeps $q_\phi(\mathbf z|\mathbf x)$ close to this prior.

- Contrast to a VAE a simple AE uses a single deterministic encoder $f_{\text{enc}}(\mathbf x)$ while VAE uses a distribution over encoders. At test time, one often uses $\pmb \mu_\phi(\mathbf x)$ as a deterministic embedding or samples $\mathbf z$ to capture uncertainty.

---

## Variational Autoencoder

### Example Architecture

![](figs/simple_vae.png){fig-cap="Simple VAE architecture (originally from Ryan D’Cunha
)" fig-alt="Simple VAE" .r-stretch fig-align="center" width=70%}

::: aside
Image originally from [Ryan D’Cunha
](https://towardsdatascience.com/author/rtdcunha/)
:::

## Takeaway

- Exact posteriors are often intractable 
    - use VI to optimise a tractable amortised posterior $q$.
- $\text{ELBO} = \mathbb E_q[\log p(\mathbf x|\mathbf z)] - \text{KL}(q(\mathbf z)\|p(\mathbf z))$ 
    - maximise it!
- VAEs are amortised approximate inference with an encoder decoder structure
    - Reparameterisation via $\mathbf z=\pmb \mu_\phi(\mathbf x)+\pmb \sigma_\phi(\mathbf x)\odot \pmb \epsilon$ with $\pmb \epsilon\sim\mathcal N(\mathbf 0,\mathbf I)$ enables backprop through sampling.
- Setting a prior $p(\mathbf z)= \mathcal N(0, \mathbf I)$ KL regularises the latent space.

---

## References

::: {#refs}
:::