Having spent a while on diffusion models I kept running into VAE's, Variational Autoencoders, in the literature. The main idea is to have a model compress down into some latent space, then poof it back up again. A very basic example is my color space VAE, which you can find under projects. It compresses the 3 main color (Red green blue) hues down into a latent 2D space. The great thing about this latent space is that you are able to sample within the latent space which results in entirely new outputs. This post however is my understanding and breakdown of the math.

What We Actually Want

The real goal is to model the distribution of our data, $p(x)$. If we knew $p(x)$, we could sample from it and generate new digits. The trick is to introduce a latent variable $z$ and assume each $x$ was produced by first drawing some hidden $z$ and then generating $x$ from it. This is comes from the Law of total probability.

$$p(x) = \int p(x \mid z)\, p(z)\, dz$$

We fix the prior $p(z) = \mathcal{N}(0, I)$ and let the decoder network learn $p(x \mid z)$. Conceptually clean, but the integral is intractable: there is no way to integrate over every possible $z$ in a high-dimensional latent space.

Introducing the Encoder

The true posterior $p(z \mid x)$,"given this image, which $z$ produced it?", is also intractable. So we approximate it with a learned Gaussian produced by the encoder:

$$q_\phi(z \mid x) = \mathcal{N}\!\left(\mu_\phi(x),\;\sigma^2_\phi(x)\right)$$

The subscript $\phi$ just means these parameters come out of the encoder network's weights. Given an input $x$, the encoder spits out a mean vector $\mu$ and a (log-)variance vector $\log \sigma^2$, and those two together define a distribution in latentspace.

The ELBO: A Tractable Lower Bound

We cannot maximise $\log p(x)$ directly, but we can derive a lower bound on it that we can maximise. Start from the definition:

$$\log p(x) = \log \int p(x, z)\, dz$$

Multiply and divide by the encoder distribution $q_\phi(z \mid x)$:

$$\log p(x) = \log \int \frac{p(x, z)}{q_\phi(z \mid x)}\, q_\phi(z \mid x)\, dz = \log\, \mathbb{E}_{q_\phi(z \mid x)}\!\left[\frac{p(x, z)}{q_\phi(z \mid x)}\right]$$

Now apply Jensen's inequality. Since $\log$ is concave, $\log \mathbb{E}[X] \geq \mathbb{E}[\log X]$:

$$\log p(x) \;\geq\; \mathbb{E}_{q_\phi(z \mid x)}\!\left[\log \frac{p(x, z)}{q_\phi(z \mid x)}\right]$$

This lower bound is called the ELBO (Evidence Lower BOund). Maximising the ELBO pushes up $\log p(x)$ — which is exactly what we want.

Splitting the ELBO Into Two Terms

Using $p(x, z) = p(x \mid z)\, p(z)$ and splitting the log:

$$\text{ELBO} = \mathbb{E}_{q_\phi(z \mid x)}[\log p(x\mid z)] + \mathbb{E}_{q_\phi(z \mid x)}\!\left[\log \frac{p(z)}{q_\phi(z \mid x)}\right]$$

The second expectation is the negative KL divergence from $q_\phi$ to the prior:

$$\text{ELBO} = \underbrace{\mathbb{E}_{q_\phi(z \mid x)}[\log p(x \mid z)]}_{\text{reconstruction}} \;-\; \underbrace{D_{KL}\!\left(q_\phi(z \mid x) \,\|\, p(z)\right)}_{\text{regulariser}}$$

Two very interpretable pieces. The reconstruction term rewards the decoder for rebuilding $x$ from samples of the encoder's distribution. The KL term pulls the encoder's distribution ( the latent space ) toward the standard normal, which is what makes the latent space usable for sampling at generation time.

The Reconstruction Term

If we model $p(x \mid z)$ as a Gaussian with fixed variance, then maximising $\log p(x \mid z)$ reduces to minimising MSE between the input and the decoder's output:

$$\log p(x \mid z) \;\propto\; -\|x - \hat{x}\|^2$$

If we instead model the pixels as Bernoulli probabilities, it becomes binary cross-entropy. Either way, this term is the familiar reconstruction loss.

The KL Divergence — Deriving the Closed Form

This is the part that took me the longest to internalise. KL divergence between two distributions is defined as:

$$D_{KL}(q \| p) = \int q(z)\, \log \frac{q(z)}{p(z)}\, dz$$

For one latent dimension with $q(z) = \mathcal{N}(\mu, \sigma^2)$ and $p(z) = \mathcal{N}(0, 1)$, the Gaussian PDFs are:

$$q(z) = \frac{1}{\sqrt{2\pi\sigma^2}} \exp\!\left(-\frac{(z-\mu)^2}{2\sigma^2}\right) \qquad p(z) = \frac{1}{\sqrt{2\pi}} \exp\!\left(-\frac{z^2}{2}\right)$$

Taking logs and subtracting, the $-\tfrac{1}{2}\log(2\pi)$ constants cancel and we are left with:

$$\log \frac{q(z)}{p(z)} = -\frac{1}{2}\log \sigma^2 - \frac{(z - \mu)^2}{2\sigma^2} + \frac{z^2}{2}$$

Now take the expectation under $q(z)$ term by term.

Term A

$\mathbb{E}_q\!\left[-\tfrac{1}{2}\log \sigma^2\right]$ is a constant with respect to $z$ — it only depends on the encoder's output — so it passes straight through:

$$\mathbb{E}_q\!\left[-\tfrac{1}{2}\log \sigma^2\right] = -\tfrac{1}{2}\log \sigma^2$$

Term B

$\mathbb{E}_q\!\left[-\tfrac{(z - \mu)^2}{2\sigma^2}\right]$. Since $z \sim \mathcal{N}(\mu, \sigma^2)$, the quantity $(z - \mu)^2$ is the squared deviation of $z$ from its own mean, and the expectation of that is, by definition, the variance:

$$\mathbb{E}_q[(z - \mu)^2] = \sigma^2 \quad\Longrightarrow\quad \mathbb{E}_q\!\left[-\frac{(z - \mu)^2}{2\sigma^2}\right] = -\frac{\sigma^2}{2\sigma^2} = -\frac{1}{2}$$

Term C (and a short detour on an identity)

$\mathbb{E}_q\!\left[\tfrac{z^2}{2}\right]$ needs $\mathbb{E}_q[z^2]$. This is where I initially got stuck, because the usual derivation jumps straight to the identity $\mathbb{E}[z^2] = \operatorname{Var}(z) + (\mathbb{E}[z])^2$ without saying where it comes from. It is worth pausing on, because it is not a trick — it is just the definition of variance rearranged.

Start from the definition of variance and let $\mu= \mathbb{E}[z]$:

$$\operatorname{Var}(z) = \mathbb{E}\!\left[(z - \mu)^2\right] = \mathbb{E}\!\left[z^2 - 2z\mu + \mu^2\right]$$

Expectation is linear, and $\mu$ is a constant (just a number — the mean), so it pulls out of any expectation it appears in:

$$\operatorname{Var}(z) = \mathbb{E}[z^2] - 2\mu\, \mathbb{E}[z] + \mu^2$$

And $\mathbb{E}[z] = \mu$ by definition, so:

$$\operatorname{Var}(z) = \mathbb{E}[z^2] - 2\mu^2 + \mu^2 = \mathbb{E}[z^2] - \mu^2$$

Rearranging gives the identity we need:

$$\mathbb{E}[z^2] = \operatorname{Var}(z) + (\mathbb{E}[z])^2$$

That is all it is — the definition of variance, expanded, solved for $\mathbb{E}[z^2]$. With this in hand, and since $z\sim q$ has $\operatorname{Var}(z) = \sigma^2$ and $\mathbb{E}[z] = \mu$:

$$\mathbb{E}_q[z^2] = \sigma^2 + \mu^2 \quad\Longrightarrow\quad \mathbb{E}_q\!\left[\frac{z^2}{2}\right] = \frac{\sigma^2 + \mu^2}{2}$$

Combining the Three Terms

Adding Terms A, B, and C:

$$D_{KL} = -\frac{1}{2}\log \sigma^2 - \frac{1}{2} + \frac{\sigma^2 + \mu^2}{2} = \frac{1}{2}\!\left(-\log \sigma^2 - 1 + \sigma^2 +\mu^2\right)$$

Or, in the form you will see quoted in most papers:

$$\boxed{D_{KL} = -\frac{1}{2}\left(1 + \log \sigma^2 - \mu^2 - \sigma^2\right)}$$

And for a $d$-dimensional latent with a diagonal covariance, each dimension contributes independently, so we just sum:

$$D_{KL} = -\frac{1}{2}\sum_{j=1}^{d}\left(1 + \log \sigma_j^2 - \mu_j^2 - \sigma_j^2\right)$$

Sanity check the intuition: if the encoder outputs$\mu = 0$ and $\sigma^2 = 1$ (i.e. it matches the prior exactly), the penalty is $-\tfrac{1}{2}(1 + 0 - 0 - 1) = 0$. Any drift awayfrom a standard normal — in either the mean or the variance — makes the penalty grow.

The Reparameterisation Trick

One more obstacle: during training, we need to sample $z \sim q_\phi(z \mid x) = \mathcal{N}(\mu, \sigma^2)$ and then backpropagate through that sampling step to updatethe encoder. But sampling is a stochastic operation — you cannot compute $\partial z / \partial \mu$ if $z$ was just drawn at random.

The fix is to push the randomness outside the learnable parameters. Rewrite the sample as a deterministic function of $\mu$, $\sigma$, and an external noise variable:

$$z = \mu + \sigma \odot \varepsilon, \qquad \varepsilon \sim \mathcal{N}(0, I)$$

Now gradients flow cleanly: $\partial z / \partial \mu = 1$ and $\partial z / \partial \sigma = \varepsilon$. The randomness $\varepsilon$ is treated as a constant input, like the data itself. This is the same change-of-variables logic that underpins the sampling form of the Gaussian — here it just gets used for gradient flow rather than for deriving a PDF.

The Final Loss

Putting it all together, for a single data point we minimise:

$$\mathcal{L} = \underbrace{\|x - \hat{x}\|^2}_{\text{reconstruction}} \;+\; \underbrace{-\frac{1}{2}\sum_{j=1}^{d}\left(1 + \log \sigma_j^2 - \mu_j^2 - \sigma_j^2\right)}_{\text{KL divergence}}$$

In practice you often weight the KL term — this isthe $\beta$-VAE variant:

$$\mathcal{L} = \text{Reconstruction} + \beta \cdot D_{KL}$$

With $\beta < 1$ you get sharper reconstructions but a messier latent space; with $\beta> 1$ the latent space becomes more disentangled but the outputs tend to blur. It is a direct lever on the reconstruction-vs-regularity tradeoff.

A practical note: the encoder outputs $\log \sigma^2$ rather than $\sigma^2$ directly. It is numerically more stable(no risk of predicting a negative variance), and it drops straight into the KL formula as written.

Training vs Generation

During training both networks areactive:

$$x \xrightarrow{\text{encoder}} (\mu, \log \sigma^2) \xrightarrow{\text{reparam}} z \xrightarrow{\text{decoder}} \hat{x}$$

At generation time, the encoder is gone entirely — we just sample from the prior and decode:

$$z \sim \mathcal{N}(0, I) \xrightarrow{\text{decoder}} \text{new image}$$

The KL term is what makes this second step work. By pulling every per-input encoder distribution toward $\mathcal{N}(0, I)$ during training, it ensures that random samples from the prior land in regions the decoder actually knows how to decode.

Conclusion

What I like about the VAE derivation is how much it hinges on elementary identities — the definition of variance, Jensen's inequality, basic properties of Gaussians. There is no single breathtaking insight; it is a chain of small, honest algebraic moves that ends somewhere genuinely useful. Next up I want to try coupling a VAE to a diffusion prior, the way Stable Diffusiondoes — which would tie this article back to the diffusion posts from earlier.

Related

VAE Machine Learning Math Gaussian Generative Models