A distribution over codes

A deterministic autoencoder maps each input to one code. Nothing tells us what codes look like in general, so we cannot pick a code at random and decode it into something plausible. The variational autoencoder fixes this by making the code a random variable with a known distribution.Kingma and Welling, Auto-Encoding Variational Bayes, ICLR 2014; independently Rezende, Mohamed and Wierstra, ICML 2014.

The model

Start from a generative story. Draw a latent code from a prior, then generate an observation from it:

$$z \sim p(z) = \mathcal{N}(0, I), \qquad x \mid z \sim p_\theta(x \mid z) = \mathcal{N}\big(h_\theta(z),\, \sigma^2 I\big).$$

The decoder $h_\theta$ is a neural network. Fitting $\theta$ by maximum likelihood needs $\log p_\theta(x) = \log \int p_\theta(x\mid z)\,p(z)\,dz$, an integral we cannot compute, along with the posterior $p_\theta(z \mid x)$. So we approximate the posterior with an encoder that outputs a Gaussian for each input:

$$q_\phi(z \mid x) = \mathcal{N}\big(\mu_\phi(x),\, \operatorname{diag}\sigma_\phi^2(x)\big).$$

The objective

For any $q$, the log-likelihood is bounded below by the evidence lower bound:

$$\log p_\theta(x) \;\ge\; \underbrace{\mathbb{E}_{q_\phi(z\mid x)}\big[\log p_\theta(x \mid z)\big]}_{\text{reconstruction}} \;-\; \underbrace{\mathrm{KL}\big(q_\phi(z\mid x)\,\big\Vert\, p(z)\big)}_{\text{stay close to the prior}} .$$

The gap is exactly $\mathrm{KL}\big(q_\phi(z\mid x)\,\Vert\, p_\theta(z\mid x)\big)$, so maximising the bound also improves the encoder’s approximation to the true posterior. With a Gaussian likelihood, the first term is a squared reconstruction error scaled by $1/2\sigma^2$, the same loss as in piece 1. The second term is a regulariser on the code, and for diagonal Gaussians it has a closed form:

$$\mathrm{KL} = \tfrac12 \sum_{i} \big(\mu_i^2 + \sigma_i^2 - \log \sigma_i^2 - 1\big).$$

To train through the sampling step, write $z = \mu_\phi(x) + \sigma_\phi(x) \odot \varepsilon$ with $\varepsilon \sim \mathcal{N}(0, I)$. This is the reparameterisation trick: the randomness moves into an input, and gradients flow through $\mu$ and $\sigma$.

Read as an autoencoder, then, a VAE is an encoder and a decoder trained on reconstruction, plus two changes: noise is injected into the code, and the code is pulled toward a standard Gaussian. Let $\sigma_\phi \to 0$ and the KL term becomes an L2 penalty $\tfrac12\lVert\mu\rVert^2$ on a deterministic code.

What it is good for, and what goes wrong

Because codes near the prior decode to plausible data, sampling $z \sim \mathcal{N}(0, I)$ and decoding generates new examples. VAEs became a standard generative model and a building block of later systems, for example as the compressor whose latent space a diffusion model works in. Their known problems: samples are often blurry under a Gaussian likelihood, and a powerful decoder can ignore the code entirely, driving the KL term to zero. This is called posterior collapse.Bowman et al., Generating Sentences from a Continuous Space, 2016, documented collapse with recurrent decoders and proposed annealing the KL term.

Holes in a plain code

The figure shows why the prior matters for generation. It trains an autoencoder with a two-number code on the dot images of piece 1, so the code can be drawn as a map. Each red point is the code of one validation image. Behind the points, the decoder’s output at each place on the map. The dashed circles are one and two spreads of $\mathcal{N}(0, I)$, and on the right are images decoded from sixteen draws of it.

With a plain code, nothing fixes the code’s scale or centre. Here the codes spread about six units along each axis, so the prior’s circles cover only a patch in the middle of the map. In a check over 200 draws, only 40% decoded to a clear dot, and those that did all sat within about half a pixel of one place, where real dots spread about six pixels.

Tick variational code. Now the KL term pulls the codes into the circles and fills them. With $\beta =$ 0.05, 62% of draws decode to a dot, spread over most of the frame, for a small cost in reconstruction: about 30% unexplained instead of 25%. Raise $\beta$ toward 1 and the KL term wins. The codes huddle at the origin, the reconstructions blur, and most draws decode to almost nothing. That is the road to posterior collapse.

The prior picks the axes, or fails to

Piece 2 showed that reconstruction alone cannot pick a basis for the code. A VAE adds a prior. Does the prior pick one?

Not if it is $\mathcal{N}(0, I)$. A standard Gaussian looks the same after any rotation. Rotate the code by $R$ and adjust the decoder by $R^\top$, and the prior term and the likelihood are both unchanged. This was the hope of disentanglement research: that a VAE’s latents, especially with the KL term up-weighted by some $\beta \gt 1$, would line up with the true factors of variation such as position, size or colour.Higgins et al., β-VAE, ICLR 2017. A large study found that, without supervision or other inductive biases, this cannot be guaranteed: for any model there are equally good ones whose latents are entangled.Locatello et al., Challenging Common Assumptions in the Unsupervised Learning of Disentangled Representations, ICML 2019. Where VAEs do find meaningful axes, part of the reason is that the diagonal posterior approximation breaks the symmetry, nudging the latents toward PCA-like directions.Rolinek, Zietlow and Martius, Variational Autoencoders Pursue PCA Directions (by Accident), CVPR 2019.

The fix suggested by piece 2 is a prior that is not rotation-invariant: a sparse one.

A Laplace prior is an L1 penalty

Replace the Gaussian prior with a Laplace prior, $p(z) \propto \exp(-\lVert z\rVert_1 / b)$, keep a linear decoder $D$ and Gaussian noise, and ask for the single most probable code, the MAP estimate:

$$\begin{aligned} \hat z &= \arg\max_z\; \log p(x \mid z) + \log p(z)\\ &= \arg\min_z\; \frac{1}{2\sigma^2}\lVert x - D z\rVert^2 + \frac{1}{b}\lVert z\rVert_1 . \end{aligned}$$

This is exactly the sparse-coding objective, with $\lambda = \sigma^2 / b$. Olshausen and Field framed sparse coding this way: a linear generative model with a sparse, heavy-tailed prior on the coefficients.Olshausen and Field, Sparse coding with an overcomplete basis set: a strategy employed by V1?, Vision Research, 1997. A non-negative code corresponds to an exponential prior on $z \ge 0$ instead.

With the Laplace prior, point $x$ at 8°, close to $d_1$. The posterior has corners on the axes, and its peak sits exactly on the $z_1$ axis: the MAP code uses one atom and sets the other to zero. Switch to the Gaussian prior. The contours become ellipses, and the peak moves off the axis and uses both atoms a little. A Gaussian prior shrinks every coefficient in proportion. A Laplace prior sets small ones exactly to zero. Turn $x$ toward the middle of the two atoms and both coefficients become non-zero under either prior. Raise the noise and the Laplace prior zeros out more.

Why interpretability uses SAEs rather than VAEs

So a sparse autoencoder’s loss is, up to the amortisation of piece 4, the MAP objective of a VAE-like model with a Laplace prior and a linear decoder. Interpretability takes the point estimate rather than the full posterior because its aim is different from a VAE’s. It is not trying to generate new activations. It wants to decompose the ones the model actually produced, faithfully and one at a time, into parts that can be read. A deterministic sparse code is the readable object, and a sparse prior is what makes its axes mean something.Fully variational sparse models exist, for example with spike-and-slab priors: Tonolini, Jensen and Murray-Smith, Variational Sparse Coding, UAI 2020. They have not been used much for interpreting language models.