Variational autoencoders: reparameterization, posterior collapse, β-VAE and VQ-VAE

Before this lesson: lesson 01 (the ELBO, its gap and amortized inference). Gradient estimators in detail: calculus, lesson 12. The rate–distortion view: information theory, lesson 10. The multivariate Gaussian density (the probability and statistics track) and backpropagation (deep learning, lesson 03). Time: about 95 minutes, plus the notebook.

---

Why this matters #

A VAE is the latent variable model of lesson 01 with a neural network in each direction, trained by stochastic gradient descent on the ELBO. Three choices make it work: the reparameterization trick (low-variance gradients through a random sample), the closed-form Gaussian KL, and the decoder likelihood, which sets the balance between the two loss terms. A wrong balance gives the classic failures: blurry samples and posterior collapse.

Today VAEs are used mostly as the first stage of larger generators. Stable Diffusion runs diffusion in the latent space of a KL-regularized autoencoder (lesson 13). VQGAN and neural audio codecs turn images and sound into discrete tokens for transformers. In science, LDCast (Leinonen et al., 2023) runs precipitation nowcasting in an autoencoder's latent space, and latent diffusion models of seismic ground motion do the same. A rule to remember: the second stage can only generate what the first-stage decoder can produce, so check the autoencoder first.

1. Predict first #

2. The picture #

The parts of a VAE #

• Encoder $q_\phi(z \mid x) = \mathcal N\big(\mu_\phi(x), \operatorname{diag}\sigma^2_\phi(x)\big)$: a network with parameters $\phi$ maps $x$ to two vectors, the mean $\mu_\phi(x)$ and the log-variance $\log\sigma^2_\phi(x)$, for a latent $z \in \mathbb R^k$.
• Decoder $p_\theta(x \mid z)$: a network with parameters $\theta$ maps $z$ to the parameters of a distribution over $x$, for example a mean $\hat x_\theta(z)$.
• Prior $p(z) = \mathcal N(0, I)$.

The loss for one data point is minus the ELBO of lesson 01, written as two terms:

The blobs tile the prior. Each data point gets a Gaussian blob in latent space. The rate pulls every blob toward $\mathcal N(0, I)$: mean to 0, width to 1. The distortion pushes blobs apart and makes them narrow, so that $z$ identifies $x$. At the balance, the blobs tile the prior. If the rate wins completely, every blob equals the prior and $z$ says nothing about $x$: posterior collapse.

The figure shows this balance exactly, for a VAE whose encoder and decoder are linear. For this model section 3 derives the optimum in closed form, so there is no training and no noise.

Interactive figure: loading…

The reparameterization trick #

The loss contains $\mathbb E_{q_\phi(z \mid x)}[\cdot]$, an average over a random sample whose distribution depends on $\phi$. We cannot backpropagate through "draw a random number". The trick moves the randomness to a fixed input:

where $\odot$ is the elementwise product. Now $z$ is a deterministic, differentiable function of $\mu$ and $\sigma$, and the noise $\varepsilon$ does not depend on $\phi$. So the gradient flows from the decoder through $z$ into the encoder. The result is the pathwise gradient estimator. The older alternative, the score-function (REINFORCE) estimator, uses only values of the loss. The figure compares them on a 1D problem: estimate $\frac{d}{d\mu}\mathbb E_{z \sim \mathcal N(\mu, 1)}[f(z)]$.

Interactive figure: loading…

3. Definitions and main results #

The Gaussian VAE objective #

For one data point $x$, with $S$ noise samples $\varepsilon_s \sim \mathcal N(0, I)$ and $z_s = \mu_\phi(x) + \sigma_\phi(x) \odot \varepsilon_s$:

It is an unbiased estimate of the ELBO, and its gradient is an unbiased estimate of the ELBO's gradient in both $\theta$ and $\phi$. $S = 1$ is standard; the minibatch averages the noise. The encoder outputs $\log\sigma^2$, an unconstrained number, and $\sigma = \exp(\frac12\log\sigma^2)$ is positive by construction. The KL is computed exactly (Theorem 1), not by sampling, which removes one source of noise.

Theorem 1: the KL divergence between two Gaussians #

Simplest case. In 1D, $\mathrm{KL}\big(\mathcal N(\mu, \sigma^2) \,\|\, \mathcal N(0, 1)\big) = \tfrac12\big(\mu^2 + \sigma^2 - 1 - \log\sigma^2\big)$.

General case. For Gaussians on $\mathbb R^d$,

For the VAE ($\mu_0 = \mu$, $\Sigma_0 = \operatorname{diag}\sigma^2$, $\mu_1 = 0$, $\Sigma_1 = I$): $\mathrm{KL} = \tfrac12\sum_{j=1}^k\big(\mu_j^2 + \sigma_j^2 - 1 - \log\sigma_j^2\big)$.

Proof of Theorem 1

1. $\log\mathcal N(z; \mu, \Sigma) = -\tfrac12\big[d\log 2\pi + \log\det\Sigma + (z - \mu)^\top\Sigma^{-1}(z - \mu)\big]$.
2. Write $\mathbb E_0$ for the expectation under $\mathcal N(\mu_0, \Sigma_0)$. Then $\mathrm{KL} = \mathbb E_0[\log\mathcal N_0 - \log\mathcal N_1] = \tfrac12\log\frac{\det\Sigma_1}{\det\Sigma_0} + \tfrac12\mathbb E_0\big[(z - \mu_1)^\top\Sigma_1^{-1}(z - \mu_1)\big] - \tfrac12\mathbb E_0\big[(z - \mu_0)^\top\Sigma_0^{-1}(z - \mu_0)\big]$.
3. Trace trick. For $z$ with mean $m$ and covariance $S$, any matrix $A$ and vector $c$: write $z - c = (z - m) + (m - c)$. The cross term has mean 0, and $\mathbb E[(z - m)^\top A(z - m)] = \operatorname{tr}\big(A\,\mathbb E[(z - m)(z - m)^\top]\big) = \operatorname{tr}(AS)$. So $\mathbb E[(z - c)^\top A(z - c)] = \operatorname{tr}(AS) + (m - c)^\top A(m - c)$.
4. With $A = \Sigma_0^{-1}$, $c = \mu_0$: the last expectation in step 2 is $\operatorname{tr}(I) = d$. With $A = \Sigma_1^{-1}$, $c = \mu_1$: the middle one is $\operatorname{tr}(\Sigma_1^{-1}\Sigma_0) + (\mu_0 - \mu_1)^\top\Sigma_1^{-1}(\mu_0 - \mu_1)$.
5. Collect the terms: this is the general formula.
6. VAE case: $\operatorname{tr}(\Sigma_0) = \sum_j\sigma_j^2$, the quadratic term is $\sum_j\mu_j^2$, and $\log\det\Sigma_1 - \log\det\Sigma_0 = -\sum_j\log\sigma_j^2$. $\blacksquare$

Each dimension costs $\tfrac12\mu_j^2 + \tfrac12(\sigma_j^2 - 1 - \log\sigma_j^2)$. Since $s - 1 - \log s \ge 0$, with equality only at $s = 1$, a dimension is free only when $\mu_j = 0$ and $\sigma_j = 1$. A variance above 1 costs nats too, not only a variance below 1.

Pathwise and score-function gradients #

We need $\nabla_\phi\,\mathbb E_{q_\phi(z)}[f(z)]$, with $f(z) = \log p_\theta(x \mid z)$. Two unbiased estimators exist:

• Score function (REINFORCE, also called the likelihood-ratio estimator): $\nabla_\phi\mathbb E_q[f] = \mathbb E_q\big[f(z)\,\nabla_\phi\log q_\phi(z)\big]$. It needs only samples of $z$, values of $f$, and $\log q_\phi$.
• Pathwise (reparameterized): with $z = g_\phi(\varepsilon)$, for example $g_\phi(\varepsilon) = \mu + \sigma \odot \varepsilon$, $\nabla_\phi\mathbb E_q[f] = \mathbb E_\varepsilon\big[(\partial z/\partial\phi)^\top\nabla_z f(z)\big]$. It needs $f$ to be differentiable in $z$ and a differentiable sampling path.

Why both are unbiased

1. Score function: $\nabla_\phi\int q_\phi(z)f(z)\,dz = \int f(z)\,\nabla_\phi q_\phi(z)\,dz = \int f(z)\,q_\phi(z)\,\nabla_\phi\log q_\phi(z)\,dz$, because $\nabla q = q\,\nabla\log q$. The last integral is $\mathbb E_q[f\,\nabla_\phi\log q_\phi]$.
2. Pathwise: $\mathbb E_{q_\phi}[f(z)] = \mathbb E_\varepsilon[f(g_\phi(\varepsilon))]$, and the distribution of $\varepsilon$ does not depend on $\phi$. So the gradient can move inside the expectation (for smooth enough $f$), and the chain rule gives $(\partial g_\phi/\partial\phi)^\top\nabla_z f$. $\blacksquare$

Why pathwise is usually better. The pathwise estimator uses the slope of $f$. The score-function estimator uses only values of $f$, multiplied by a random factor $\nabla_\phi\log q$ with mean 0. Add a constant $c$ to $f$: the mean of the score-function estimator does not change (because $\mathbb E_q[\nabla_\phi\log q] = 0$), but its variance grows like $c^2$ times the Fisher information. The noise also grows with the dimension of $z$. A VAE's $f$ is a sum over thousands of pixels with a large magnitude, so the score-function estimate would be very noisy. Example (the figure): for $f(z) = z^2$ and $z \sim \mathcal N(2, 1)$, the gradient in $\mu$ is 4, and the per-sample variance is 4 (pathwise) against 87 (score function). "Usually", not always: for a fast-changing $f$ such as $\sin(kz)$, the pathwise variance grows like $k^2$, and the score function can win (variance about 4.5 against 0.5 for $k = 3$).

When the pathwise estimator is not available.

1. Discrete $z$: there is no differentiable path from noise to a sample. Use the score function with a baseline, a continuous relaxation (Gumbel-softmax, biased), or the straight-through estimator (biased; VQ-VAE below).
2. $f$ not differentiable in $z$, for example a step function or a black-box simulator. Then $f' = 0$ almost everywhere and the "pathwise" average is wrong, as in the figure.
3. Distributions without a simple sampling path (Gamma, Beta, Dirichlet): these need implicit reparameterization.

The decoder variance and "MSE + β·KL" #

Take a Gaussian decoder $p_\theta(x \mid z) = \mathcal N(\hat x_\theta(z), \sigma^2 I)$ for $x \in \mathbb R^n$ ($n$ pixels or features):

Proposition 2 (a fixed variance is a weight). Multiply $-\mathcal L$ by $2\sigma^2$: $2\sigma^2 \cdot (-\mathcal L) = \mathbb E_q\|x - \hat x\|^2 + 2\sigma^2\,\mathrm{KL} + \text{const}$. So minimizing "summed squared error $+\,\beta \cdot \mathrm{KL}$" is the same as maximizing the ELBO of a Gaussian decoder with

The weight $\beta$ says how noisy the model believes the data are. A large $\sigma^2$ makes it cheap to give up reconstruction (toward collapse). A small $\sigma^2$ makes reconstruction dominate (toward a plain autoencoder with an irregular latent space).

Proposition 3 (a learned variance). If one scalar $\sigma^2$ is learned with the rest, its optimum is the mean squared error per dimension, $\sigma_*^2 = \mathbb E_q\|x - \hat x\|^2/n$. Proof. Set the derivative in $\sigma^2$ to zero: $-\mathbb E_q\|x - \hat x\|^2/(2\sigma^4) + n/(2\sigma^2) = 0$. $\blacksquare$ At the optimum the distortion is $\frac n2\log(2\pi e \cdot \mathrm{MSE})$: logarithmic in the MSE. In effect $\beta = 2\,\mathrm{MSE}$, which shrinks as the reconstructions improve. A decoder that outputs a variance per pixel does the same per dimension.

Blurriness. For squared error, the best decoder mean at $z$ is the average of all inputs that the encoder maps near $z$. An average of several images is blurry. Papers usually show the decoder mean $\hat x(z)$; adding the decoder noise $\sigma \cdot$ (noise) would give grainy images instead.

Posterior collapse #

Definition. Posterior collapse: $q_\phi(z \mid x) \approx p(z)$ for every $x$, in some or all latent dimensions. The decoder ignores those dimensions, and their KL is about 0.

Proposition 4 (collapse can be optimal). Suppose some decoder ignores $z$ and still fits the data: $p_{\theta^*}(x \mid z) = p_{\text{data}}(x)$ for every $z$. Then $q(z \mid x) = p(z)$ gives a global maximum of the expected ELBO.

Proof of Proposition 4

1. The decoder ignores $z$, so $z$ and $x$ are independent under the model, and the true posterior equals the prior.
2. So $q = p(z)$ has zero gap (lesson 01, Theorem 1), and $\mathrm{ELBO}(x) = \log p_{\theta^*}(x) = \log p_{\text{data}}(x)$.
3. For any model, $\mathbb E_{\text{data}}[\mathrm{ELBO}] \le \mathbb E_{\text{data}}[\log p_\theta(x)] \le \mathbb E_{\text{data}}[\log p_{\text{data}}(x)]$; the second inequality is Gibbs' inequality. Step 2 reaches this ceiling. $\blacksquare$

Information in $z$ costs rate; the model pays it only when the reconstruction gain is larger. Autoregressive decoders (PixelCNN, LSTM or transformer text decoders) can model $p(x)$ alone, so they often collapse (Bowman et al., 2016, for text).

A linear example: collapse without a powerful decoder (Lucas, Tucker, Grosse and Norouzi, Don't Blame the ELBO!, 2019). Data $x \sim \mathcal N(0, \lambda)$ in 1D; decoder $p(x \mid z) = \mathcal N(wz, \sigma^2)$ with a fixed $\sigma^2$; prior $\mathcal N(0, 1)$; encoder $\mathcal N(ax, v)$. Train with weight $\beta$ on the KL ($\beta = 1$ is the ELBO). Then the optimum is

So the direction is active iff $\lambda > \beta\sigma^2$. If $\beta\sigma^2 \ge \lambda$, collapse is the global optimum: the decoder noise already explains all the data variance, and there is no powerful decoder at all.

Derivation of the linear optimum

1. Expected distortion: $\mathbb E_x\mathbb E_q[(x - wz)^2] = \lambda(1 - wa)^2 + w^2 v$, so $\mathbb E[\log p(x \mid z)] = -\tfrac12\log(2\pi\sigma^2) - \frac{\lambda(1 - wa)^2 + w^2 v}{2\sigma^2}$.
2. Expected rate (Theorem 1): $\tfrac12(a^2\lambda + v - 1 - \log v)$.
3. The objective is $J = \mathbb E[\log p(x \mid z)] - \beta \cdot \text{rate}$. Write $c = \beta\sigma^2$. Setting $\partial J/\partial v = 0$ gives $v = c/(c + w^2)$; setting $\partial J/\partial a = 0$ gives $a = w/(w^2 + c)$.
4. Put these back, with $u = w^2$: $J = -\tfrac12\log(2\pi\sigma^2) - \frac{1}{2\sigma^2}\Big[\frac{\lambda c}{u + c} + c\log\frac{u + c}{c}\Big]$.
5. $\frac{dJ}{du} = \frac{c\,(\lambda - u - c)}{2\sigma^2(u + c)^2}$: positive for $u < \lambda - c$, negative after. So $u_* = \lambda - c$ if $\lambda > c$, else $u_* = 0$.
6. At an active optimum, $v = c/\lambda$ and $a^2\lambda = (\lambda - c)/\lambda$, so the rate is $\tfrac12\big(\frac{\lambda - c}{\lambda} + \frac c\lambda - 1 - \log\frac c\lambda\big) = \tfrac12\log\frac{\lambda}{c}$. $\blacksquare$

In more dimensions (probabilistic PCA with data variances $\lambda_1 \ge \lambda_2 \ge \dots$ along the principal directions), each direction is treated separately: direction $i$ is active iff $\lambda_i > \beta\sigma^2$. This is the figure in section 2. So partial collapse is not always a bug: it can be the model choosing its latent dimension.

When collapse happens. (i) Strong decoders that model $p(x)$ without $z$. (ii) A large $\beta$, or a large fixed decoder variance (as in the linear example, and in question 1 of section 1). (iii) Early training: while the decoder is weak, $q \approx p(z)$ is an attractive state, and the encoder lags behind (He et al., 2019).

How to detect it. Log the KL per latent dimension, $\mathbb E_x\,\mathrm{KL}(q(z_j \mid x) \,\|\, p(z_j))$; a value near 0 marks a collapsed dimension. Count active units: dimension $j$ is active if $A_j = \operatorname{Var}_x\big(\mathbb E_q[z_j \mid x]\big) > 0.01$ (Burda et al., 2016). Check that decoded outputs change when you change $z$. Report rate and distortion separately, not only their sum.

Fixes. KL annealing: raise the KL weight from 0 to 1 during training (also cyclical schedules). Free bits (Kingma et al., 2016): replace $\sum_j\mathrm{KL}_j$ by $\sum_j\max(\lambda, \mathrm{KL}_j)$, with $\mathrm{KL}_j$ averaged over the minibatch, so the first $\lambda$ nats per dimension are free. Learn the decoder variance, so that $\sigma^2$ can drop below the data variance. Weaken the decoder, or add skip connections from $z$ to every decoder layer. More encoder updates per decoder update (He et al., 2019) help only when the encoder lags behind (cause iii), not when collapse is the optimum.

The β-VAE: the rate–distortion view #

Average over the data: distortion $D = \mathbb E_x\mathbb E_q[-\log p_\theta(x \mid z)]$ and rate $R = \mathbb E_x\,\mathrm{KL}(q_\phi(z \mid x) \,\|\, p(z))$. Then $-\mathbb E_x[\mathcal L] = D + R$. The β-VAE (Higgins et al., 2017) minimizes

• $\beta > 1$: lower rate, more distortion; more compressed and often more factorized codes, blurrier reconstructions. Since $R \ge 0$, $-(D + \beta R) \le -(D + R) \le \mathbb E_x\log p_\theta(x)$: still a bound, but looser.
• $\beta < 1$: more information in $z$, sharper reconstructions, a weaker match between the codes and the prior. $-(D + \beta R)$ is in general not a bound on $\mathbb E_x\log p_\theta(x)$.
• The rate bounds the information that $z$ carries about $x$: $R \ge I_q(x; z)$ (information theory, lesson 10).

Predicting the choice. Each trained model is a point $(R, D)$. The objective $D + \beta R$ picks the point that a line of slope $-\beta$ touches first when it moves up from below. So the winner always lies on the lower convex hull of the points. Two models $A$, $B$ with $R_A < R_B$ tie at $\beta^* = (D_A - D_B)/(R_B - R_A)$; above $\beta^*$ the lower-rate model wins.

Interactive figure: loading…

VQ-VAE #

A VQ-VAE (vector-quantized VAE; van den Oord, Vinyals and Kavukcuoglu, 2017) uses a discrete latent: a list of vectors instead of a Gaussian blob.

• Codebook. $K$ vectors $e_1, \dots, e_K \in \mathbb R^d$. The encoder outputs $z_e(x) \in \mathbb R^d$ (one per spatial position in practice).
• Quantization. $k^* = \arg\min_j\|z_e(x) - e_j\|^2$ and $z_q = e_{k^*}$: the nearest code. So $q(z = k \mid x)$ is one-hot. With a uniform prior over the $K$ codes, the KL term is the constant $\log K$ and is dropped.
• Loss (with $\operatorname{sg}$ = stop-gradient: the identity in the forward pass, zero gradient in the backward pass):

with $\beta = 0.25$ in the paper (results were similar for $0.1$ to $2.0$). This $\beta$ is not the β of the β-VAE.

• Straight-through gradient. In the forward pass, the decoder sees $z_q$. In the backward pass, $\partial\mathcal L/\partial z_q$ is copied unchanged to $z_e$, as if quantization were the identity. In code: $z_q = z_e + \operatorname{sg}[e_{k^*} - z_e]$.
• Who trains what. The decoder: the reconstruction term. The encoder: reconstruction (through the straight-through copy) and the commitment term, which keeps $z_e$ close to its code so that encoder outputs do not drift or grow. The codebook: only the codebook term, which moves $e_{k^*}$ toward the encoder outputs assigned to it (often replaced by moving averages: an online k-means).
• Generation. The prior is uniform during training. Afterwards, an autoregressive model (PixelCNN, a transformer; lesson 03) is fitted on the code indices, sampled, and decoded.
• Codebook collapse. Often only a few codes are used. An unused code is never the argmin, so it gets no gradient and stays dead. The usage $u_k$ (the fraction of positions that pick code $k$) gives the perplexity $\exp(-\sum_k u_k\log u_k)$, the effective number of codes. Fixes: initialize codes from encoder outputs (k-means), moving-average updates, reset dead codes to recent encoder outputs, low-dimensional $\ell_2$-normalized codes, or finite scalar quantization (FSQ; Mentzer et al., 2023), which rounds a few dimensions to a fixed grid and has no learned codebook to collapse.

Interactive figure: loading…

Why autoencoders are the first stage. An autoencoder trained with reconstruction, perceptual and adversarial losses (VQGAN; Esser et al., 2021) removes detail that the eye does not see and keeps a small, dense latent. In the Stable Diffusion autoencoder, a $256 \times 256 \times 3$ image (196,608 numbers) becomes a $32 \times 32 \times 4$ latent (4,096 numbers, 48 times fewer), and the KL weight is only about $10^{-6}$ (lesson 13). The diffusion model or the transformer then spends its capacity on content, and the cost of attention falls with the number of tokens. Discrete codes serve token models (VQGAN + transformer, residual VQ in audio codecs); continuous, weakly regularized latents serve diffusion.

Answers to section 1. (1) Several hundred. Multiply the loss by $n = 784$: it becomes summed squared error $+\ 784 \cdot \mathrm{KL}$. So $\beta = 784$ and, by Proposition 2, $\sigma^2 = \beta/2 = 392$ ($\sigma \approx 19.8$ for pixels in $[0, 1]$). The model believes that the pixels are almost pure noise, so ignoring $z$ is cheap: collapse and a blurry average image. Fix: sum the squared error over pixels ($\sigma^2 = \tfrac12$) or learn $\sigma^2$. (2) Straight-through: the encoder receives $\partial\mathcal L/\partial z_q$ as if quantization were the identity, plus the gradient of the commitment term.

4. Worked example #

(a) Gaussian KL. The encoder outputs $\mu = (2, 0)$ and $\sigma = (1, 0.5)$, so $\sigma^2 = (1, 0.25)$.

• Dimension 1: $\tfrac12(\mu_1^2 + \sigma_1^2 - 1 - \log\sigma_1^2) = \tfrac12(4 + 1 - 1 - 0) = 2$.
• Dimension 2: $\tfrac12(0 + 0.25 - 1 - \log 0.25) = \tfrac12(-0.75 + 1.3863) = 0.3181$.
• Total: $2.3181$ nats. Dimension 1 pays for its mean, dimension 2 for a variance below 1.

(b) One reparameterized sample. With $\varepsilon = (0.5, -1)$: $z = (2 + 1 \cdot 0.5,\ 0 + 0.5 \cdot (-1)) = (2.5, -0.5)$. The local derivatives are $\partial z_j/\partial\mu_j = 1$ and, since $\sigma = \exp(\tfrac12\log\sigma^2)$, $\partial z_j/\partial\log\sigma_j^2 = \tfrac12\sigma_j\varepsilon_j = (0.25, -0.25)$. Backpropagation multiplies $\nabla_z f$ by these numbers.

(c) The decoder variance.

• Summed squared error $+$ KL: $\beta = 1$, so $\sigma^2 = \tfrac12$.
• Summed squared error $+\ 0.02 \cdot$ KL: $\sigma^2 = 0.01$ ($\sigma = 0.1$).
• Mean squared error over $n = 784$ pixels $+$ KL: $\beta = 784$, $\sigma^2 = 392$.
• A learned scalar variance with a per-pixel MSE of $0.01$: $\sigma_*^2 = 0.01$.

(d) Collapse in the linear VAE. Data variance $\lambda = 1$, $\beta = 1$.

• $\sigma^2 = 0.25$: $w_*^2 = 1 - 0.25 = 0.75$, so $w_* = 0.866$; rate $= \tfrac12\log(1/0.25) = \tfrac12\log 4 = 0.693$ nats. The encoder family contains the exact posterior, so the ELBO equals the log-likelihood of $\mathcal N(0, w_*^2 + \sigma^2) = \mathcal N(0, 1)$: $-\tfrac12\log(2\pi) - \tfrac12 = -1.4189$.
• $\sigma^2 = 1.5$: $\lambda < \sigma^2$, so $w_* = 0$ and the rate is 0 (collapse). The model is $\mathcal N(0, 1.5)$, and ELBO $=$ log-likelihood $= -\tfrac12\log(3\pi) - \frac{1}{2 \cdot 1.5} = -1.4550$.

(e) Which model does the β-VAE prefer? Three trained models, in nats:

model	$R$	$D$	$D + 0.2R$	$D + R$	$D + 2R$
A	2	60	60.4	62	64
B	10	50	52.0	60	70
C	25	44	49.0	69	94

Ties: A and B at $\beta^* = (60 - 50)/(10 - 2) = 1.25$; B and C at $\beta^* = (50 - 44)/(25 - 10) = 0.4$. So C wins for $\beta < 0.4$, B for $0.4 < \beta < 1.25$ (this includes the ELBO, $\beta = 1$), and A for $\beta > 1.25$.

(f) One VQ step. Codebook $e_1 = (0, 0)$, $e_2 = (2, 0)$, $e_3 = (0, 2)$; encoder output $z_e = (1.5, 0.4)$.

• Squared distances: $\|z_e - e_1\|^2 = 2.25 + 0.16 = 2.41$, $\|z_e - e_2\|^2 = 0.25 + 0.16 = 0.41$, $\|z_e - e_3\|^2 = 2.25 + 2.56 = 4.81$. So $k^* = 2$ and $z_q = (2, 0)$.
• Codebook term: $\|z_e - e_2\|^2 = 0.41$ (summed over the 2 entries); commitment term: $0.25 \times 0.41 = 0.1025$.
• Gradients: the codebook term gives $\nabla_{e_2} = 2(e_2 - z_e) = (1, -0.8)$, so a descent step moves $e_2$ toward $z_e$. The commitment term gives $\nabla_{z_e} = 2\beta(z_e - e_2) = (-0.25, 0.2)$, which pulls $z_e$ toward $e_2$. The reconstruction gradient at $z_q$ is added to $\nabla_{z_e}$ unchanged. $e_1$ and $e_3$ get no gradient.

5. In code #

NumPy and SciPy check the numbers of section 4 by other routes: Monte Carlo for the KL, per-sample gradient estimates, a numerical optimizer for the linear VAE, and SciPy's vector quantizer.

import numpy as np
from scipy.stats import norm
from scipy.optimize import minimize
from scipy.cluster.vq import vq

# (a) Closed-form KL(N(mu, diag s^2) || N(0, I)) vs a Monte Carlo estimate
mu, s = np.array([2.0, 0.0]), np.array([1.0, 0.5])
kl = 0.5 * (mu**2 + s**2 - 1 - np.log(s**2))
rng = np.random.default_rng(0)
z = mu + s * rng.normal(size=(1_000_000, 2))                    # reparameterized samples
mc = (norm.logpdf(z, mu, s) - norm.logpdf(z, 0, 1)).mean(0)
print(kl.round(4), kl.sum().round(4), mc.round(3))
# [2.     0.3181] 2.3181 [2.004 0.318]

# (b) Two unbiased per-sample estimates of d/dmu E[z^2], z ~ N(2, 1); the true value is 4
eps = rng.normal(size=1_000_000)
pathwise = 2 * (2.0 + eps)                                     # f'(mu + eps)
score = (2.0 + eps) ** 2 * eps                                 # f(z) * d/dmu log q(z) = f(z) (z - mu)
print(f"{pathwise.mean():.2f} {pathwise.var():.1f} | {score.mean():.2f} {score.var():.0f}")
# 4.00 4.0 | 3.99 86   (exact variances: 4 and 87)

# (c) Linear VAE: x ~ N(0, lam), decoder N(w z, s2) with s2 fixed, encoder N(a x, v).
#     The expected ELBO in closed form, maximized over (w, a, log v).
def elbo_and_rate(params, lam, s2):
    w, a, logv = params
    v = np.exp(logv)
    recon = -0.5 * np.log(2 * np.pi * s2) - (lam * (1 - w * a) ** 2 + w**2 * v) / (2 * s2)
    rate = 0.5 * (a**2 * lam + v - 1 - logv)
    return recon - rate, rate

lam = 1.0
for s2 in (0.25, 1.5):
    res = minimize(lambda p: -elbo_and_rate(p, lam, s2)[0], x0=[0.5, 0.5, 0.0])
    elbo, rate = elbo_and_rate(res.x, lam, s2)
    var_x = res.x[0] ** 2 + s2                                 # model marginal: x ~ N(0, w^2 + s2)
    loglik = -0.5 * np.log(2 * np.pi * var_x) - lam / (2 * var_x)
    print(s2, f"|w| = {abs(res.x[0]):.3f}  rate = {rate:.3f}  ELBO = {elbo:.4f}  log-lik = {loglik:.4f}")
# 0.25 |w| = 0.866  rate = 0.693  ELBO = -1.4189  log-lik = -1.4189
# 1.5 |w| = 0.000  rate = 0.000  ELBO = -1.4550  log-lik = -1.4550   (collapse is optimal)

# (d) Which (rate, distortion) model does D + beta R prefer?
R = np.array([2.0, 10.0, 25.0]); D = np.array([60.0, 50.0, 44.0])
for beta in (0.2, 1.0, 2.0):
    print(beta, D + beta * R, "ABC"[np.argmin(D + beta * R)])
# 0.2 [60.4 52.  49. ] C | 1.0 [62. 60. 69.] B | 2.0 [64. 70. 94.] A

# (e) Nearest code of the VQ example (SciPy's vector quantization)
codebook = np.array([[0.0, 0.0], [2.0, 0.0], [0.0, 2.0]])
print(vq(np.array([[1.5, 0.4]]), codebook))
# (array([1], dtype=int32), array([0.64031242])): index 1 is e_2 (0-based), distance^2 = 0.41

The score-function variance is itself a noisy estimate with heavy tails: with $10^6$ samples it lands near 87, not exactly on it. In PyTorch, torch.distributions.Normal(mu, sigma).rsample() gives the reparameterized sample and .sample() gives a sample with no gradient path.

6. Where it shows up #

• Latent diffusion (lesson 13). The Stable Diffusion autoencoder adds perceptual (LPIPS) and patch-adversarial losses to a very small KL term. In LDCast (precipitation) and in latent models of seismic ground motion, check that the autoencoder reconstructs the quantities you will evaluate (extremes, spectra, peak values) before you train the diffusion stage.
• Token-based generators. VQGAN codes plus a transformer or a masked model; DALL·E 1 used a discrete VAE trained with a Gumbel-softmax relaxation; audio codecs (SoundStream, EnCodec) use residual VQ under speech and music language models.
• Representation learning. β-VAE and its variants are used to learn more factorized ("disentangled") codes; the linear analysis shows that β also decides how many directions survive.
• Training logs. Plot rate, distortion, per-dimension KL and active units. A rate near 0 with a good-looking loss is collapse, not success.

7. Common mistakes #

• Mean-reduced reconstruction with a summed KL. This is a silent $\beta = n$, the number of pixels; see question 1 of section 1.
• An encoder that outputs $\sigma$ with no positivity constraint. Output $\log\sigma^2$ (and clamp it in practice) to avoid NaNs.
• Binary cross-entropy on non-binary pixels. For $x \in [0, 1]$ it is not a normalized likelihood, so the "ELBO" is not a bound on a real $\log p(x)$ and cannot be compared with other models.
• Reading a small KL as good regularization. Check the active units first.
• Wrong stop-gradients in VQ. Writing $z_q = e_{k^*}$ gives the encoder no reconstruction gradient; dropping $\operatorname{sg}$ in the codebook term lets it train the encoder too.
• Reporting a $\beta < 1$ objective as a likelihood bound. Only $\beta \ge 1$ guarantees a bound.
• Using the pathwise gradient through a discrete or non-differentiable step. The gradient is 0 almost everywhere, so the estimate is silently wrong.

8. Check yourself #

1. The encoder outputs $\mu = (0, 0, 1)$ and $\sigma^2 = (1, 4, 1)$. Compute the KL to $\mathcal N(0, I)$.
   Answer

   Dimension 1: 0. Dimension 2: $\frac12(4 - 1 - \log 4) = 0.8069$. Dimension 3: $\frac12 \cdot 1 = 0.5$. Total $1.3069$ nats.
2. A paper minimizes $\|x - \hat x\|^2$ (summed) $+\ 0.02 \cdot \mathrm{KL}$. Which decoder does this assume?
   Answer

   A Gaussian decoder with fixed variance $\sigma^2 = \beta/2 = 0.01$ ($\sigma = 0.1$), by Proposition 2.
3. Why can the reparameterization trick not be used for a categorical latent? Name two alternatives.
   Answer

   There is no differentiable map from fixed noise to a discrete sample: the sample changes in jumps, and its derivative with respect to the class probabilities is 0 almost everywhere. Alternatives: the score-function estimator with a baseline; a Gumbel-softmax relaxation; the straight-through estimator.
4. Linear VAE with data variance $\lambda = 2$, fixed decoder variance $\sigma^2 = 0.25$ and $\beta = 1$: give $w_*$ and the rate. What changes for $\sigma^2 = 3$?
   Answer

   $w_*^2 = 2 - 0.25 = 1.75$, $w_* = 1.323$; rate $\frac12\log(2/0.25) = \frac12\log 8 = 1.040$ nats. For $\sigma^2 = 3 \ge \lambda$: $w_* = 0$ and the rate is 0; posterior collapse is the optimum.
5. Models P ($R = 4$, $D = 30$) and Q ($R = 12$, $D = 24$). Which wins for $\beta = 1$, for $\beta = 0.5$, and where do they tie?
   Answer

   $\beta = 1$: P gives 34, Q gives 36, so P. $\beta = 0.5$: P gives 32, Q gives 30, so Q. Tie at $\beta^* = (30 - 24)/(12 - 4) = 0.75$.
6. A VQ-VAE has 512 codes; after training, 40 codes are used, about equally often. What is the effective number of bits per position, why do the other 472 codes not recover, and what would you do?
   Answer

   About $\log_2 40 = 5.32$ bits instead of $\log_2 512 = 9$. A dead code is never the nearest code, so the codebook term never updates it. Reset dead codes to recent encoder outputs, initialize with k-means, use moving-average updates or normalized low-dimensional codes, or switch to FSQ.

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

1. Give $\mathrm{KL}(\mathcal N(\mu, \sigma^2) \,\|\, \mathcal N(0, 1))$ in 1D.

   Show answer

2. Write the reparameterization trick for a diagonal Gaussian encoder.

   Show answer

3. Write the score-function (REINFORCE) gradient estimator.

   Show answer

4. Why does the pathwise estimator usually have lower variance than the score-function estimator?

   Show answer

5. Which decoder variance does "summed squared error $+\ \beta \cdot$ KL" assume?

   Show answer

6. Define rate and distortion of a VAE.

   Show answer

7. What is posterior collapse, and how do you detect it?

   Show answer

8. In a linear VAE with fixed decoder variance $\sigma^2$, when is a latent direction with data variance $\lambda$ active?

   Show answer

9. What is the straight-through estimator in VQ-VAE?

   Show answer

10. Name the three loss terms of VQ-VAE.

    Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "Write the VAE loss." — Summed reconstruction NLL plus $-\frac12\sum_j(1 + \log\sigma_j^2 - \mu_j^2 - \sigma_j^2)$, with $z = \mu + \sigma \odot \varepsilon$. Watch the reductions (sum vs mean).
• "Why does reparameterization reduce variance?" — It uses $\nabla_z f$; the score-function estimator uses only values of $f$ times a zero-mean factor, so its noise scales with $|f|$.
• "Is posterior collapse an optimization bug?" — Not always: with a decoder that can model $p(x)$ alone, or a decoder variance above the data variance, it is the optimum of the ELBO.
• "What does β do in a β-VAE? Is the objective still a bound?" — It trades distortion for rate along the rate–distortion curve; it is a bound only for $\beta \ge 1$.
• "How do gradients pass the argmin in VQ-VAE? What is the commitment loss for?" — A straight-through copy; the commitment loss keeps encoder outputs close to their codes.
• "Why are VAE samples blurry and diffusion samples sharp?" — A Gaussian decoder mean averages over what $z$ does not determine. A diffusion sampler builds the output in many small denoising steps, so the final output is a sample, not an average over the remaining uncertainty.

10. Practice #

The notebook for this lesson runs in Colab. It has five exercises, each with a check cell and a folded solution, and an optional PyTorch demo:

1. The Gaussian KL in closed form vs Monte Carlo, in 1D and in the general (full-covariance) case.
2. The variance of pathwise and score-function gradients for three test functions.
3. A linear VAE in 2D: optimize it numerically and find the β at which each latent direction collapses; compare with $\lambda_i/\sigma^2$.
4. Nearest-code quantization and the straight-through gradient: check which gradient reaches the encoder and the codebook.
5. Choose between trained models with $D + \beta R$ and find the tie values of β.
6. (Optional, PyTorch) a tiny VAE on two moons: the 2D latent space, samples, and active units.

11. Watch and read #

Variational AutoencodersArxiv Insights · 15 min

The encoder, the decoder and the reparameterization trick, with pictures of the latent space.

Variational Autoencoders | Generative AI AnimatedDeepia · 20 min

An animated derivation of the VAE loss from the ELBO.

Vector-Quantized Variational Autoencoders (VQ-VAEs)DeepBean · 18 min

The codebook, the straight-through gradient and the commitment loss.

• Diederik P. Kingma and Max Welling, An Introduction to Variational Autoencoders (2019): §2.3–2.4 (gradients and the reparameterization trick), §2.5 (the factorized Gaussian encoder), §2.8 (KL annealing, free bits and blurriness).
• Aaron van den Oord, Oriol Vinyals and Koray Kavukcuoglu, Neural Discrete Representation Learning (VQ-VAE, 2017): §3.1–3.3 (the loss, $\beta = 0.25$, the prior fitted afterwards).
• James Lucas, George Tucker, Roger Grosse and Mohammad Norouzi, Don't Blame the ELBO! A Linear VAE Perspective on Posterior Collapse (2019): the linear analysis behind the figures of this lesson.
• Irina Higgins et al., β-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework (ICLR 2017).
• Lilian Weng, From Autoencoder to Beta-VAE: a compact overview with the reparameterization trick, β-VAE and VQ-VAE.
• Kevin Murphy, Probabilistic Machine Learning: Advanced Topics (free PDF): ch. 21 (VAEs), including §21.4 (posterior collapse) and §21.6 (VQ-VAE).