A trained VAE goes wrong in two familiar ways. Its samples come out blurry, like smudged averages of training images. And when its decoder is strong, for example one that writes a sentence one word at a time, it can learn to ignore its code completely. The "autoencoder" has then become an ordinary language model (Bowman et al. 2016). Neither failure is a bug. In both cases the optimizer does exactly what the training objective asks.

This section shows that both failures come from one trade-off inside the VAE's loss. The loss charges for every piece of information the code carries, and it also charges for every error the decoder makes. One weight, $\beta$, sets the exchange rate between the two charges. Once you see the trade-off you can predict which details a VAE drops first, measure whether its code is being used, and pick the fix that matches the cause. When a VAE stops using its code, the failure is called posterior collapse. "Why are VAE samples blurry?" and "what is posterior collapse?" are among the most common interview questions on generative models, and this section answers both precisely.

Here is the setup, recapped from §32.3. A data point $\mathbf{x} \in \R^{d_x}$, such as an image, is described by a short code $\mathbf{z} \in \R^{d_z}$. The encoder maps an input to a Gaussian distribution over codes, $q_\phi(\mathbf{z}\mid\mathbf{x}) = \mathcal{N}\big(\boldsymbol\mu_\phi(\mathbf{x}), \operatorname{diag}\boldsymbol\sigma^2_\phi(\mathbf{x})\big)$, where $\phi$ stands for the encoder's weights. We call this distribution the encoder's posterior, because it plays the part of the posterior of Bayes rule, the distribution of the code once the input is known. The decoder $p_\theta(\mathbf{x}\mid\mathbf{z})$, with weights $\theta$, maps a code back to a distribution over inputs. The prior $p(\mathbf{z}) = \mathcal{N}(\mathbf{0}, \mathbf{I})$ is where codes come from when we generate. The prior and the decoder together also define an exact posterior $p_\theta(\mathbf{z}\mid\mathbf{x})$ by Bayes rule, which we call the true posterior to keep it apart from the encoder's. Training maximizes the evidence lower bound of §32.2, averaged over the training set,

In plain words, the first term rewards reconstructing $\mathbf{x}$ from a code drawn from the encoder. The second term is the KL divergence, which measures how far the encoder's distribution for this input sits from the prior. It is measured in nats, the unit of information when logarithms are natural (1 nat is $1/\ln 2 = 1.44$ bits). The whole objective is a lower bound on the log-likelihood $\log p_\theta(\mathbf{x})$.

The KL term is what this section is about. It charges the model for every nat of information that the code carries about the input. Sometimes a blurred guess costs less than the detail would. Sometimes the decoder can do the whole job without the code. In both cases the optimizer takes the cheaper option. The plan is to make that exact. We name the two terms of the loss rate and distortion, and show that a weight $\beta$ on the KL term picks one point on a trade-off curve between them. A linear VAE can be solved in closed form. It meets the best limit that information theory allows, and collapse shows up in it as directions that are switched off. That picture then explains posterior collapse and its fixes, splits the KL term into a useful part and a wasted part, and shows what $\beta$-VAEs and hierarchical VAEs can and cannot do.

Two formulas from §32.3 recur throughout. A Gaussian decoder with mean network $\mathbf{f}_\theta$ and a fixed variance $\sigma_x^2$, $p_\theta(\mathbf{x}\mid\mathbf{z}) = \mathcal{N}(\mathbf{f}_\theta(\mathbf{z}), \sigma_x^2\mathbf{I})$, has negative log-likelihood

a squared error scaled by $1/(2\sigma_x^2)$ plus a constant. And the KL of a one-dimensional Gaussian $\mathcal{N}(\mu, \sigma^2)$ to the standard normal is $\tfrac12(\mu^2 + \sigma^2 - 1 - \ln\sigma^2)$ nats.

Why VAE samples are blurry

There are two separate reasons, and it pays to keep them apart. The first is about the decoder's mean, the image it outputs for a code. The second is about the shape of the decoder's distribution around that mean.

The best decoder averages the images that share a code

Suppose two different training images could each have produced the same code. A decoder trained with squared error cannot tell which one it is looking at. Its best answer is then the average of the two, weighted by how likely each one is to have produced that code. An average of two images is a blend, and a blend of shifted images is a blur. Here is the smallest case with numbers.

The same argument works for any number of images, and it is the conditional-mean theorem of §32.1 applied to the encoder's codes. That theorem says the function that minimizes an expected squared error is a conditional mean, whatever the input is.

In plain words, the decoder's output at a code is the average of every training image that could have produced that code, weighted by how likely each one is. When the posteriors of different images overlap, every code in the overlap decodes to a blend. And the KL term pushes every posterior toward $\mathcal{N}(\mathbf{0},\mathbf{I})$, which means wide and overlapping. Narrow, well-separated posteriors are exactly what the KL charges for. Rezende and Viola (2018) make the same point about overlapping posteriors and blurred reconstructions.

Two images show the mechanism but not the blur we see in real samples. Blur needs many images that differ by small shifts, so that the images sharing a code are neighbors. So take 16 one-dimensional "images" of 16 pixels each, numbered $c = 0, 1, \dots, 15$. Image $c$ is a narrow bar, a smooth bump of brightness centered on pixel $c$, and neighboring images differ by a one-pixel shift. The code is one number. The encoder sends image $c$ to $\mathcal{N}(a\,\bar c, s^2)$, where $\bar c = c - 7.5$ is the bump's position measured from the middle. Two numbers set the encoder, the spacing $a$ between the code means and the width $s$ of each posterior.

What decides the blur is the ratio $r = s/a$, the width of one posterior in units of the spacing. With $r$ small the posteriors do not touch, and with $r$ large they pile on top of each other. The rate, the average KL to the prior, depends on both numbers. For a given overlap $r$ the cheapest spacing is $a^2 = 1/(v + r^2)$, where $v = 21.25$ is the variance of the positions $\bar c$, and the rate is then

In plain words, the rate grows as the posteriors narrow relative to the spread of their means. Information theory knows this formula as the capacity of a Gaussian channel whose signal-to-noise ratio is $v/r^2$.

The cell computes the best decoder on a grid of codes. For each $\beta$ it picks the overlap that minimizes the loss, the squared error over $2\sigma_x^2$ plus $\beta$ times the rate, with $\sigma_x^2 = 0.1$. It prints the rate, the squared error, and how many images the decoder effectively averages at an image's own code. That last number is $e^{H(\mathbf{w})}$, the exponential of the entropy of the weights $w_c$. It is 1 when one image gets all the weight and 16 when all sixteen get equal weight. The cell also plots the decoded bar.

So blur and collapse are not two different failures. They are two positions on one curve, and the rest of the section is about that curve.

In practice, this means a VAE's blur gets worse when you raise the KL weight and better when you let the code spend more nats. An interviewer who asks "why are VAE samples blurry" wants this argument first.

A factorized decoder cannot say "here or there"

The second reason is about the decoder's distribution, not its mean. Take a Gaussian decoder with a diagonal covariance, $p_\theta(\mathbf{x}\mid\mathbf{z}) = \prod_i \mathcal{N}(x_i; f_i(\mathbf{z}), \sigma_x^2)$, where $f_i(\mathbf{z})$ is its mean for pixel $i$. It treats whatever the code leaves undecided as independent noise on each pixel. Suppose the code fixes "a face, looking left" but not the exact position of one eyelash. The true distribution of images given that code is then a set of sharp images, each with the eyelash in a slightly different place. No product of independent per-pixel Gaussians can represent that set. The best one, in the sense of the highest expected log-likelihood, puts its mean at the pixelwise average, a soft gray streak. It spreads the rest as per-pixel noise.

Sampling from that decoder gives the gray streak plus white speckle, which looks worse than the mean. So papers display the mean $\mathbf{f}_\theta(\mathbf{z})$ instead, and that displayed mean is the blur. Some decoders can represent correlated uncertainty, such as autoregressive pixel models that generate one pixel at a time given the ones before it. We will see shortly that they cause the opposite problem, a decoder so good that it ignores the code. The practical answers are adversarial and perceptual losses that let the decoder commit to one sharp texture (§32.6), and discrete codes (§32.5).

One more setting acts on both reasons, the decoder variance $\sigma_x^2$. The squared error is divided by $2\sigma_x^2$ and the KL is multiplied by $\beta$, so raising $\sigma_x^2$ does the same thing as raising $\beta$. §32.3 called this a hidden KL weight, and the linear model below shows that only the product $\beta\sigma_x^2$ matters. So a VAE trained with a large fixed $\sigma_x^2$ is blurry for the same reason as one trained with a large $\beta$. Its posteriors overlap more, which is the first reason. Its code also carries less, so more of each image is left undecided for the factorized decoder to average, which is the second.

Rate, distortion, and beta as a price

The two terms of the negative ELBO measure two different costs, and they deserve names. The rate is how many nats the code spends describing each example. The distortion is how many nats the decoder still needs to describe the example once it has the code. The names come from information theory, where a rate counts the bits sent per symbol and a distortion is the error that the description leaves. Alemi et al. (2018) introduced this reading of the VAE.

The weight $\beta$ is a price. Minimizing $\mathcal{D} + \beta\mathcal{R}$ buys a nat of rate only if it lowers the distortion by at least $\beta$ nats. With $\beta = 2$, for example, an extra nat in the code must save 2 nats of distortion to be worth it. A large $\beta$ makes rate expensive, so the optimum has low rate and high distortion.

There is a picture for this. Consider every encoder and decoder the architecture can represent, and plot the $(\mathcal{R}, \mathcal{D})$ pair of each one. The lower-left edge of that cloud is a curve $\mathcal{D}(\mathcal{R})$, the least distortion achievable at each rate. Minimizing $\mathcal{D} + \beta\mathcal{R}$ lands on that edge, at the point where a straight line of slope $-\beta$ just touches it. Where the curve is smooth, setting the derivative of $\mathcal{D}(\mathcal{R}) + \beta\mathcal{R}$ to zero gives $\mathrm{d}\mathcal{D}/\mathrm{d}\mathcal{R} = -\beta$. In plain words, each $\beta$ selects one point on the trade-off curve, and $\beta$ is the slope there.

What point does the ordinary $\beta = 1$ choose? When the decoder is powerful enough to model the data by itself, it may choose to send nothing at all.

In plain words, if the decoder can do the job alone, then at $\beta = 1$ using the code can at best tie with ignoring it. It ties only if the encoder matches the true posterior exactly. With any approximation error in the encoder, using the code loses. Above $\beta = 1$ ignoring the code wins outright. Below $\beta = 1$ rate is cheap, and the model prefers to use the code. The cases that made collapse famous are discrete data, sentences in Bowman et al. (2016) and binarized MNIST in Alemi et al. (2018), so the theorem applies to them as stated.

The theorem shows up in real models. On binarized MNIST, a VAE trained just above $\beta = 1$ ignored its code, and one trained at $\beta = 1$ used it, with almost the same ELBO.

The ELBO alone cannot tell a model that uses its code from one that ignores it. In practice, log the rate and the distortion separately, never only their sum. "Two VAEs have the same ELBO, so are they equally good?" is a classic interview trap, and the answer is no.

The best possible trade-off for Gaussian data

To judge any of these trade-offs we need the limit, the least rate that any description could get away with at a given error. Information theory supplies it. The mutual information $I(\mathbf{x};\hat{\mathbf{x}})$ between a source $\mathbf{x}$ and a reconstruction $\hat{\mathbf{x}}$ is how many nats knowing one tells you about the other (§0.5). Shannon's rate-distortion function is the least mutual information that any randomized description can have while keeping the expected squared error at most $\mathcal{D}_{\text{sq}}$,

Shannon's rate-distortion theorem says what this number means in practice. It is the least number of nats per source vector that any code needs to reach that error, when many source vectors are described together (Cover and Thomas, chapter 10).

Start with one number. Let $x$ be Gaussian with variance 4, and allow a squared error of 1. The answer is $\tfrac12\ln(4/1) = 0.693$ nats, exactly 1 bit. In general one Gaussian number with variance $\lambda$ costs $\tfrac12\ln(\lambda/D)$ nats at error $D$, so each extra bit divides the error variance by 4. With several independent numbers the question becomes how to split the error budget among them, and the answer has a picture.

Draw each variance $\lambda_i$ as a bar and pour water to height $\nu$. Wherever the water covers a bar, that component is not described at all. Wherever a bar sticks out, the part above the water is what the description carries, and the water underneath is the error left behind. The rule is to describe the large components and drop the small ones. A component gets any rate only once the error on every larger component has been pushed below that component's own variance.

Once you know which components sit under the water, the water level has a closed form. Computing it is a short and common coding question.

import numpy as np

def water_fill(eigs, D):
    lam = np.asarray(eigs, dtype=float)
    nu = D / len(lam)                     # placeholder: assumes every component is above the water
    Di = np.full_like(lam, nu)
    return nu, float(np.sum(0.5 * np.log(lam / nu))), Di

The linear beta-VAE is reverse water-filling

The simplest VAE whose optimum we can find exactly is the linear one. Here we ask where $\beta$ puts it. Given enough latent dimensions, it lands exactly on Shannon's limit, and $\beta$ sets the water level.

The data are centered, with covariance $\boldsymbol\Sigma_x = \mathbf{U}\boldsymbol\Lambda\mathbf{U}^\top$. Its eigenvalues are $\lambda_1 \ge \dots \ge \lambda_{d_x}$ and its orthonormal eigenvectors $\mathbf{u}_i$ are the principal directions of §0.10, the directions along which the data vary most. The decoder is linear, $p_\theta(\mathbf{x}\mid\mathbf{z}) = \mathcal{N}(\mathbf{W}\mathbf{z}, \sigma_x^2\mathbf{I})$ with $\mathbf{W}\in\R^{d_x\times d_z}$ and a fixed $\sigma_x^2$. The encoder is linear too, with a diagonal covariance shared by all inputs, $q_\phi(\mathbf{z}\mid\mathbf{x}) = \mathcal{N}(\mathbf{V}^\top\mathbf{x}, \operatorname{diag}(\mathbf{v}))$ with $\mathbf{V}\in\R^{d_x\times d_z}$. This is the linear VAE of §32.3, also used by Lucas et al. (2019). We drop the constant $\tfrac{d_x}{2}\log(2\pi\sigma_x^2)$ from the distortion and write $\mathcal{D}' = \mathcal{D}_{\text{sq}}/(2\sigma_x^2)$.

Recall that probabilistic PCA (pPCA) is the model $\mathbf{x} = \mathbf{W}\mathbf{z} + \text{noise}$ with $\mathbf{z} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$ and Gaussian noise of a fixed variance. Its maximum-likelihood $\mathbf{W}$ points along the top principal directions, and it switches off every direction whose variance is below the noise (§32.2). Its exact posterior has the same covariance for every input, which is why our encoder needs no input-dependent variance.

In plain words, $\beta$ and the decoder variance matter only through their product, the water level. Every principal direction whose variance clears the water is encoded, and it leaves exactly $\nu$ of squared error behind. Every direction below the water is switched off, and its latent dimension sits exactly on the prior. That switched-off dimension is posterior collapse in its cleanest form, and it is not a training failure. It is the optimum. Lucas et al. (2019) made this point for $\beta = 1$, where the water level is the decoder variance itself.

Here are the numbers for the chapter's running spectrum.

A theorem about the global optimum is only useful if gradient descent finds it. The cell trains a linear $\beta$-VAE on a five-dimensional spectrum $(4, 2, 1, 0.5, 0.25)$, rotated by a random orthogonal matrix so that the principal directions are not the coordinate axes. It uses $\sigma_x^2 = 0.5$ and $\beta = 1.5$, so $\nu = 0.75$, and $d_z = 4$. It descends the exact expected loss, whose gradients are short matrix expressions.

Gradient descent lands on the closed form to four decimals. Three directions clear the water at 0.75 and get rates of 0.837, 0.490, and 0.144 nats. The fourth latent dimension carries nothing. It is free to take the direction with eigenvalue 0.5, but 0.5 is below the water, so encoding that direction would cost more rate than it saves in error. The squared error is $3 \times 0.75 + 0.5 + 0.25 = 3.0$. The next cell traces the whole curve for the running spectrum as $\beta$ varies, for each latent size.

In the plot, the solid black curve is Shannon's limit for this spectrum, and the dashed curve is the linear VAE with $d_z = 2$. That VAE follows the limit exactly until the water drops below the third eigenvalue. After that the third direction has no latent dimension to go to, and the curve leaves the limit. Each $\beta$ is one point on the curve, and the slope there is $-1/(2\beta\sigma_x^2)$, steeper for small $\beta$.

The widget below draws the same curves, under the eigenvalues drawn as bars with the water level across them. It fixes the decoder variance at $\sigma_x^2 = 0.5$, since only the product $\beta\sigma_x^2$ matters. Change the spectrum, $\beta$, and the number of latent dimensions $d_z$, or press Raise β, and watch directions sink below the water one at a time. Try the flat spectrum, where every direction crosses the water together and the code collapses all at once.

This is the cleanest answer to an interview question that comes up often, "which latent dimensions of a VAE collapse, and why?" In the linear model the answer is exact. The directions whose variance is below $\beta\sigma_x^2$ collapse, and it is optimal that they do.

import numpy as np

def linear_beta_vae(eigs, sigma2_x, beta, d_z):
    raise NotImplementedError

Why the blur is a low-pass filter

The linear theorem also says which details a VAE drops first, the directions of smallest variance. For images that has a concrete meaning. Call an image model stationary if its statistics do not depend on position, with the image wrapped around at its borders. The eigenvectors of a stationary model's covariance are the sines and cosines of the Fourier basis. So its principal directions are spatial frequencies, patterns of stripes at a given spacing and angle. Its eigenvalues form the power spectrum, the variance at each frequency. The power spectra of natural images fall off roughly as $1/f^2$ with spatial frequency $f$, an empirical regularity measured across many photographs (Field 1987; van der Schaaf and van Hateren 1996). Fine detail, the edges and texture, lives at high frequencies and has the smallest variance.

So reverse water-filling spends any budget on the lowest frequencies and drops everything above a cutoff. An image with its high frequencies removed is, by definition, a blurred image. The cell builds a $64\times64$ image model with a $1/f^2$ spectrum and finds the water level for several budgets in bits per pixel. It then shows a random image from the model next to its optimal reconstruction at 0.05 bits per pixel.

At 0.05 bits per pixel the optimal code keeps about 7 percent of the Fourier modes, all of them below a cutoff of about 9.5 cycles per image. Its reconstruction explains just over half the variance. Each kept mode is also shrunk by $1 - \nu/\lambda_i$, the posterior-mean factor from the derivation. The reconstruction looks like the image seen through frosted glass. Real VAEs are nonlinear and images are not Gaussian, so this is a model, not a measurement of any particular VAE. But the pressure it isolates is real. At a finite rate the ELBO saves most by dropping the details that carry the least variance, and in natural images those are the fine ones.

Posterior collapse

We have now seen collapse twice, as the gray corner of the bar toy and as directions below the water. Here is the definition, and how to measure it.

The two statistics answer different questions. Activity asks whether the code changes with the input, which is how a Gaussian encoder carries information in practice. The per-dimension rate asks what the dimension costs, and a dimension can cost something without carrying anything. Take an encoder whose output in dimension $j$ is $\mathcal{N}(3, 1)$ for every input. It pays $\tfrac12(9 + 1 - 1 - 0) = 4.5$ nats per example and tells the decoder nothing. Later in the section we will see exactly where such wasted nats live.

Both statistics take a few lines of code, and computing them is the first thing to do when a VAE looks suspicious.

import numpy as np

def latent_usage(mu, logvar, thresh=1e-2):
    raise NotImplementedError

Collapse has two causes

Collapse has two different causes, and they call for different fixes. Telling them apart is the most useful skill in this section.

Collapse at the optimum. Sometimes collapse is what the objective wants. The linear theorem showed one case, where directions below the water level switch off at the global optimum. The theorem on powerful decoders showed another. Chen et al. (2017) state the general principle, which they call information preference. Any information that the decoder can model by itself is cheaper to model there. Routing it through the code costs an extra $D_{\mathrm{KL}}(q_\phi(\mathbf{z}\mid\mathbf{x})\Vert p_\theta(\mathbf{z}\mid\mathbf{x}))$, the gap between the ELBO and the log-likelihood (§32.2), while modeling it in the decoder costs nothing extra. An autoregressive decoder predicts each token or pixel from the previous ones, which is exactly what a language model does (§6.2). It can model nearly everything by itself, so it takes nearly everything.

Bowman et al. (2016) saw this on the Penn Treebank, a standard corpus of newspaper sentences. A VAE with an LSTM decoder, a recurrent network that writes one word at a time, was trained plainly. It converged to zero KL and the same likelihood as the LSTM language model alone. With two fixes defined below, KL annealing and word dropout, their reported model kept a KL of 2 nats. Its test negative log-likelihood bound was 101 nats per sentence, against the language model's exact 100.

Collapse by optimization. Other times the objective would like to use the code, but training never finds that solution. He et al. (2019) traced the mechanism. At initialization the encoder's $q_\phi(\mathbf{z}\mid\mathbf{x})$ and the model's true posterior $p_\theta(\mathbf{z}\mid\mathbf{x})$ are both close to the prior. As the decoder learns, its true posterior starts to move. The encoder chases this moving target and lags behind, so the codes the decoder sees say little about the input. The gap term $D_{\mathrm{KL}}(q_\phi(\mathbf{z}\mid\mathbf{x})\Vert p_\theta(\mathbf{z}\mid\mathbf{x}))$ then pulls the true posterior back toward the uninformative encoder. The decoder learns to ignore the code, and both settle at the prior. Kingma et al. (2016) describe the same trap. Near the start, $q_\phi \approx p$ is attractive, the encoder's gradients are mostly noise, and the state is a stable equilibrium that is hard to leave.

The fixes, and which cause each one treats

Each fix below changes either the objective or only the path of training. A fix that changes the objective can act on collapse at the optimum. A fix that changes only the path can act only on collapse by optimization. Knowing which is which predicts when a fix will work.

• KL annealing, or warm-up (Bowman et al. 2016; Sønderby et al. 2016). Replace $\beta$ by a schedule $\beta_t$ that rises from 0 to 1 over training, linearly as $\beta_t = \min(1, t/T_{\text{warm}})$ at step $t$, or along a sigmoid. Early on the code is free, so the decoder learns to use it before the KL is charged in full. It changes only the path, since the final objective is the ordinary ELBO.
• Cyclical annealing (Fu et al. 2019). Split training into $M$ cycles of $T_{\text{cyc}}$ steps and ramp within each cycle, $\beta_t = \min(1, \tau_t/\rho)$, where $\tau_t = (t \bmod T_{\text{cyc}})/T_{\text{cyc}}$ is the position inside the cycle and $\rho$ is the fraction of each cycle spent ramping. The paper's defaults are $M = 4$ and $\rho = 0.5$. Each return to $\beta = 0$ lets the decoder learn again from informative codes, starting from the better codes of the previous cycle.
• Free bits (Kingma et al. 2016), defined in the box below. A floor on each group's rate, below which the KL is not charged. This changes the objective.
• A committed rate (Razavi et al. 2019). Choose the posterior and prior families so that their KL can never fall below some $\delta \gt 0$, which the paper calls a $\delta$-VAE. The simplest version gives the posterior a fixed variance $\sigma_q^2 \neq 1$. Then $D_{\mathrm{KL}}(\mathcal{N}(\mu, \sigma_q^2)\Vert\mathcal{N}(0,1)) \ge \tfrac12(\sigma_q^2 - 1 - \ln\sigma_q^2) \gt 0$ whatever the mean. The paper's sequence models instead pair an independent posterior with a prior that is correlated in time. Either way the minimum rate is paid whether or not the code is used, so collapsing no longer saves it.
• Weaker decoders. Word dropout replaces a random fraction of the decoder's input words with an unknown-word token, so the decoder must get that information from the code (Bowman et al. 2016). The variational lossy autoencoder restricts an autoregressive image decoder to a small receptive field, so the code must carry the global structure (Chen et al. 2017).
• Aggressive inference (He et al. 2019). Before each decoder update, update the encoder repeatedly until the ELBO stops improving, so it never lags. In their runs that took about 30 to 100 encoder steps per decoder step. They do this only early in training, until the mutual information between inputs and codes stops rising, and then return to ordinary updates.
• Target the rate directly. Burgess et al. (2018) replace $\beta\mathcal{R}$ by $\gamma\lvert\mathcal{R} - C\rvert$, with a large $\gamma$ (1,000 in their experiments) and a target capacity $C$ raised slowly during training (to 25 nats on the dSprites shapes dataset). GECO, short for generalized ELBO with constrained optimization (Rezende and Viola 2018), instead minimizes the KL subject to a constraint on the reconstruction error. The constraint enters the loss with a Lagrange multiplier, a weight that training raises while the error is above its target and lowers once it is below.

In plain words, each group gets $\lambda_{\text{fb}}$ nats of rate for free, averaged over the batch. While a group spends less than that, its KL term has zero gradient. The only force acting on it is then the reconstruction term, which rewards using the code. When the posterior and prior factorize across groups, as diagonal Gaussians do, the free-bits penalty is never smaller than the true KL. So the objective is still a lower bound on the log-likelihood, just a looser one.

Why the batch average? The floor is meant for a rate, and a rate is an average over data. Clamping each example separately would switch the KL gradient off example by example. It would hand every example its own allowance, however much the group carries on average. And it would loosen the bound further, since the average of a maximum is at least the maximum of the average.

Which of these fixes help in the linear model, where collapse is at the optimum? The cell trains the running example three ways, with $\sigma_x^2 = 0.5$, $\beta = 1$, and $d_z = 3$, so the third direction, $\lambda = 0.25$, sits below the water at 0.5. It trains plainly, with a linear warm-up over the first half of training, and with free bits at 0.25 nats per dimension.

In deep VAEs, where collapse is often an optimization trap, annealing does help. That is why NVAE warms up its KL term, and why the Ladder VAE paper found warm-up crucial for keeping its upper layers active (Sønderby et al. 2016). Both models appear in the hierarchical VAEs below. The table collects the fixes.

In training code the schedules and the floor combine into one function of the step.

import numpy as np

def kl_penalty(kl, step, schedule="constant", warmup=1000, cycle=1000, ratio=0.5, free_bits=0.0):
    raise NotImplementedError

The same experiment with a small nonlinear VAE in PyTorch shows the same behaviors. The data have two true factors pushed through a nonlinear map into 12 dimensions. The model has 8 latent dimensions, and a decoder variance of 0.3 sets a high water level.

The plain model uses 3 of its 8 dimensions and leaves the rest on the prior. Warm-up ends in the same place. Free bits lift every dimension to at least the 0.1-nat floor. The weakest of the three used dimensions drops from 0.5 to 0.1 nats, because its information is spread over the newly free ones, the same respreading the linear model showed. The activity statistic now counts all 8 units as active, even though six of them carry only about the free 0.1 nats. So a jump in active units after adding free bits is not, by itself, evidence of a better code. Whether the respreading is a fix or a waste depends on what you wanted the code for.

What the KL term really pays for

The rate pays for two different things, and only one of them helps. Splitting it in two explains both the wasted nats of the constant-offset encoder and why samples from the prior can look bad. The split needs one new object. The aggregate posterior $q_\phi(\mathbf{z}) = \frac1N\sum_n q_\phi(\mathbf{z}\mid\mathbf{x}_n)$ is the distribution of codes over the whole dataset, the cloud you get by encoding every training example.

The decomposition is due to Hoffman and Johnson (2016), who called it ELBO surgery. In plain words, the rate pays for information, how much the code says about which example it came from. That can never exceed $\log N$ nats, the cost of naming one example out of $N$. The rate also pays for a mismatch, how far the cloud of all codes is from the prior. Only the information helps reconstruction. The mismatch is pure cost, and it is also what breaks sampling. A VAE generates by drawing $\mathbf{z}$ from $p(\mathbf{z})$, but the decoder has only been trained on codes from $q_\phi(\mathbf{z})$. Regions where the prior has mass and the aggregate posterior has little are called prior holes. Codes drawn there decode to things the decoder was never asked to produce, often blends like the halfway codes in the bar widget.

The bar toy makes this exact, because its code is one number and the integrals can be done on a grid.

With wide, overlapping posteriors almost all of the rate is information. As they narrow, the information climbs to its ceiling, $\ln 16 = 2.773$ nats, and stops. Past that point every extra nat of rate buys nothing for reconstruction and goes entirely into the mismatch, since sixteen narrow spikes look less and less like a standard normal. That is the precise sense in which a VAE trained to low distortion on a finite dataset gets bad samples.

Four responses follow from the decomposition.

• Choose a better prior. For a fixed encoder, the prior that minimizes the mismatch is the aggregate posterior itself. The VampPrior approximates it with a mixture of the encoder's posteriors at $m$ learned pseudo-inputs $\mathbf{a}_1, \dots, \mathbf{a}_m$, points in data space, $p(\mathbf{z}) = \frac1m\sum_{k=1}^{m} q_\phi(\mathbf{z}\mid\mathbf{a}_k)$ (Tomczak and Welling 2018).
• Fit a second model to the codes. The two-stage VAE trains a second, small VAE on samples from $q_\phi(\mathbf{z})$ and generates through both (Dai and Wipf 2019). Latent diffusion takes the same idea to its conclusion, fitting a diffusion model to the codes of an autoencoder.
• Split further. The mismatch itself splits into a total correlation, $D_{\mathrm{KL}}(q_\phi(\mathbf{z})\Vert\prod_j q_\phi(z_j))$, which measures dependence between latent dimensions, and the sum of per-dimension mismatches (Chen et al. 2018). Penalizing the total correlation gives a handle on disentanglement, which a later subsection takes up.
• Stop charging for information. Penalize only the mismatch, which is the family of models in the next subsection.

import numpy as np

def elbo_surgery(mu, logvar, S, rng):
    mean_kl = float((0.5 * (mu**2 + np.exp(logvar) - 1.0 - logvar)).sum(1).mean())
    return mean_kl, 0.0, mean_kl          # placeholder: pretends the code carries no information

Autoencoders that charge only for the mismatch

The decomposition says the rate pays for two things. Only one of them, the mismatch $D_{\mathrm{KL}}(q_\phi(\mathbf{z})\Vert p(\mathbf{z}))$, is what breaks sampling. The other, the information $I_q$, is what pushes the code to say less. Charging for it is what lets ignoring the code tie or win when the decoder is strong. A family of autoencoders keeps the first charge and drops the second. They minimize

where $c$ is a reconstruction cost (squared error, or a negative log-likelihood), $\mathbb{D}$ is any divergence that is zero exactly when its two arguments are equal, and $\lambda_{\text{agg}} \gt 0$ is a weight. There is no per-example KL. In plain words, reconstruct well, and make the cloud of all codes look like the prior, but let each individual code say as much about its input as it likes.

The difficulty is that $q_\phi(\mathbf{z})$ is a mixture over the whole dataset, so $\mathbb{D}$ has no closed form. It can be estimated from a minibatch, though. The codes of a minibatch, $\mathbf{z}_b \sim q_\phi(\mathbf{z}\mid\mathbf{x}_b)$ with $\mathbf{x}_b$ drawn from the data, are samples from $q_\phi(\mathbf{z})$, and samples from $p(\mathbf{z})$ cost nothing. Three models differ only in how they estimate $\mathbb{D}$ from those two batches.

• The adversarial autoencoder (AAE, Makhzani et al. 2015) trains a discriminator in code space to tell codes from prior samples, and trains the encoder to fool it. In its ideal form that game minimizes the Jensen-Shannon divergence between $q_\phi(\mathbf{z})$ and $p(\mathbf{z})$, as §32.6 derives for the image-space game of a tokenizer. The encoder can be deterministic, because the discriminator only ever sees samples.
• The MMD-VAE of InfoVAE (Zhao, Song, and Ermon 2017) uses the maximum mean discrepancy, a kernel distance between two samples defined below.
• The Wasserstein autoencoder (WAE, Tolstikhin et al. 2018) arrives at the same objective from optimal transport, and estimates $\mathbb{D}$ either with a discriminator (WAE-GAN) or with the maximum mean discrepancy (WAE-MMD).

The maximum mean discrepancy is the batch estimator these models made standard. Writing its unbiased form is a common interview coding question.

In plain words, the codes should be as similar to prior samples, on average, as each set is to itself. The estimate leaves out the diagonal terms $k(\mathbf{z}_i, \mathbf{z}_i)$, because a point compared with itself is not a draw of two independent points. Because it is unbiased, it is often slightly negative when $Q = P$, which is correct, since its average is zero. §33.7 uses the same estimator to score generated samples against a known distribution.

Which kernel? The WAE authors found that the RBF kernel fails exactly when it is needed most. Early in training some codes land far from the prior. There $\exp(-\lVert\mathbf{z} - \tilde{\mathbf{z}}\rVert^2/(2\ell^2))$ and its gradient are numerically zero, so nothing pulls those codes back. The IMQ kernel decays only as the inverse square of the distance, so far codes still feel a pull. WAE used it with $C = 2d_z\sigma_z^2$ for a prior $\mathcal{N}(\mathbf{0}, \sigma_z^2\mathbf{I})$, which is the expected squared distance between two prior samples.

What do these models buy, and what do they give up? With no charge for information, reconstruction rewards every nat the code carries. So the incentive to collapse is gone even with a powerful decoder, and the mismatch penalty still shapes the code space for sampling. The cost is that $\mathcal{L}_{\text{agg}}$ is not a bound on any likelihood, so these are autoencoders with a well-shaped code space, not density models. And nothing limits how narrow each posterior becomes, up to the $\log N$ ceiling on a finite dataset. InfoVAE reports that keeping the encoder's variances from vanishing is enough in practice, and WAE accepts the extreme case and uses deterministic encoders.

import numpy as np

def mmd_penalty(Zq, Zp, kernel_scale, kernel="imq"):
    raise NotImplementedError

Beta-VAE and the limits of disentanglement

So far a large $\beta$ has been the problem. A $\beta$-VAE sets $\beta$ above 1 on purpose, to make the code expensive (Higgins et al. 2017). The authors trained on synthetic images generated from a few independent factors, such as position, scale, rotation, and shape. With an expensive code, single latent dimensions often lined up with single factors, so that moving one latent number moved one factor. This property is called disentanglement. The hope was that it would make representations easier to interpret and make downstream tasks need fewer labeled examples.

Why should an expensive code be disentangled? The rate-distortion picture gives the most honest answer. Raising $\beta$ raises the water level, so the code keeps only the few directions of largest variance. With a diagonal posterior, each latent dimension prefers to hold one direction. In the linear model this is exact, since the optimum aligns latent dimensions with the principal directions whenever the eigenvalues are distinct. Rolinek et al. (2019) argue that the diagonal posterior has the same effect in deep VAEs. But principal directions are directions of large variance, and nothing guarantees that they are the factors a human would name. Burgess et al. (2018) give a related account, in which a slowly raised capacity $C$ lets factors enter the code one at a time, the ones that most reduce the reconstruction error first.

There is a deeper problem. The standard normal prior looks the same in every rotated coordinate system, and the likelihood cannot break that symmetry. For any orthogonal matrix $\mathbf{Q}$, the rotated code $\mathbf{z}' = \mathbf{Q}\mathbf{z}$ has the same distribution $\mathcal{N}(\mathbf{0},\mathbf{I})$. The decoder $\mathbf{z}'\mapsto\mathbf{f}_\theta(\mathbf{Q}^\top\mathbf{z}')$ then defines exactly the same distribution over data. So the likelihood gives no reason to prefer one set of latent axes over any rotation of them. Locatello et al. (2019) showed that this problem is not special to Gaussians.

In plain words, take any generative model with independent factors and a decoder. Composing the decoder with $\mathbf{f}^{-1}$ gives a second model with the same prior and the same distribution over data, whose latent coordinates each mix all the original factors. A method that sees only data cannot tell the two apart. So if it finds the disentangled one, that was due to its inductive biases, the assumptions built into the architecture, the diagonal posterior, and the optimizer, not to the data.

Interpretability research meets the same wall from another side. The sparse autoencoders of §23.6 get their features from an explicit sparsity bias, not from the data alone. An interviewer may ask whether a $\beta$-VAE "discovers the true factors". The precise answer is that it finds high-variance directions aligned with its diagonal posterior. Nothing in the data alone can single out the true factors.

Hierarchical VAEs

The most successful answer to blur and collapse in likelihood-based VAEs is depth in the latent space itself. A hierarchical VAE has $L$ groups of latents generated top-down, each group conditioned on the ones above it,

where $\mathbf{z}_{\gt l} = (\mathbf{z}_{l+1}, \dots, \mathbf{z}_L)$. The top groups are small and coarse, such as a $1\times1$ or $4\times4$ map in an image model, and the lower groups are fine. The encoder uses the same top-down order, $q_\phi(\mathbf{z}_{1:L}\mid\mathbf{x}) = \prod_l q_\phi(\mathbf{z}_l\mid\mathbf{x}, \mathbf{z}_{\gt l})$. The chain rule of KL divergence then splits the rate into one term per group,

It follows by writing the log ratio of the two products as a sum of per-group log ratios and taking the expectation one group at a time. In plain words, each group's prior is now learned and conditional, so a group pays only for what the groups above it did not already predict. Collapse can happen group by group, and the per-group rate is the thing to monitor.

Three designs mark the progress.

• Ladder VAE (Sønderby et al. 2016) introduced the top-down encoder. Each group's posterior merges a bottom-up Gaussian estimate $\mathcal{N}(\hat{\boldsymbol\mu}_l, \hat{\boldsymbol\sigma}^2_l)$, computed from $\mathbf{x}$, with the top-down prior $\mathcal{N}(\boldsymbol\mu_{p,l}, \boldsymbol\sigma^2_{p,l})$. The merge is a precision-weighted average, $\boldsymbol\sigma_{q,l}^{-2} = \hat{\boldsymbol\sigma}_l^{-2} + \boldsymbol\sigma_{p,l}^{-2}$ and $\boldsymbol\mu_{q,l} = \boldsymbol\sigma_{q,l}^2\odot(\hat{\boldsymbol\mu}_l\odot\hat{\boldsymbol\sigma}_l^{-2} + \boldsymbol\mu_{p,l}\odot\boldsymbol\sigma_{p,l}^{-2})$, which is the product of the two Gaussian densities. The authors found warm-up and batch normalization crucial for keeping the upper groups active.
• NVAE (Vahdat and Kautz 2020) writes each posterior relative to its prior, $q = \mathcal{N}(\mu_p + \Delta\mu, (\sigma_p\Delta\sigma)^2)$ in each dimension. The KL becomes $\tfrac12\big(\Delta\mu^2/\sigma_p^2 + \Delta\sigma^2 - \ln\Delta\sigma^2 - 1\big)$, so it depends mainly on the encoder's relative outputs $\Delta\mu$ and $\Delta\sigma$. NVAE adds spectral regularization, a penalty on each layer's largest singular value, to keep the KL from blowing up. During KL warm-up it also balances the KL across groups so that none switches off. It reached 2.91 bits per dimension on CIFAR-10, in the units of §32.3.
• VDVAE (Child 2021) went deep, with 45 stochastic layers on CIFAR-10 for 2.87 bits per dimension, and needed neither free bits nor KL warm-up. Its appendix shows the architecture deciding collapse. In a 73-layer model on ImageNet $32\times32$, removing the residual connections from the top-down blocks left the first 45 layers encoding nothing. A learned upsampler left the first 13 encoding nothing. With nearest-neighbor upsampling, the first layers all encoded information.

Two formulas from these designs are worth being able to write from memory, the Ladder merge and NVAE's residual KL. The chain rule above is the check that ties them together. In a two-level model whose every piece is linear and Gaussian, the joint posterior and the joint prior are both Gaussians over $(\mathbf{z}_1, \mathbf{z}_2)$. So the joint KL can be computed directly with the full-covariance Gaussian KL of §32.3 and compared with the sum of the per-group KLs.

import numpy as np

def ladder_merge(mu_hat, logvar_hat, mu_p, logvar_p):
    raise NotImplementedError

def nvae_residual_kl(delta_mu, log_delta_sigma, sigma_p):
    raise NotImplementedError

Two neighbors of this idea deserve a sentence each. A richer posterior can replace depth. The inverse autoregressive flow pushes a diagonal Gaussian through invertible autoregressive maps so that the posterior can be correlated (Kingma et al. 2016). That is a normalizing flow, a model family that flow matching revisits in continuous time. And a diffusion model is a hierarchical VAE with two special choices. Its encoder is fixed rather than learned, with each layer adding a little Gaussian noise, and it has as many layers as noise levels. That is how diffusion models derive their training loss from the same ELBO.

Choosing the KL weight in practice

The ELBO cannot choose a point on the rate-distortion curve for you (Alemi et al. 2018), so practitioners choose one deliberately, by what the code is for. Two production systems show how far from 1 that choice can land.

The second choice does not contradict anything above. A latent-diffusion autoencoder is never sampled from its prior. A diffusion model is trained afterwards on its codes, which is the two-stage idea, so prior holes do not matter and there is no reason to pay for matching $\mathcal{N}(\mathbf{0},\mathbf{I})$. The choice comes down to what the code is for.

Discrete codes are where §32.5 begins.

Read more

• Alemi et al. (2018). Fixing a Broken ELBO. The rate-distortion view of VAEs, the bounds $H - \mathcal{D} \le I \le \mathcal{R}$, and experiments showing equal ELBOs with very different codes.
• Hoffman & Johnson (2016). ELBO surgery: yet another way to carve up the variational evidence lower bound. Four pages that split the KL term into mutual information and marginal KL, with the MNIST measurements quoted above (archived copy, since the workshop site is gone).
• Lucas et al. (2019). Don't Blame the ELBO! A Linear VAE Perspective on Posterior Collapse. The linear VAE, its correspondence with pPCA, and why the observation noise decides which directions collapse.
• He et al. (2019). Lagging Inference Networks and Posterior Collapse in Variational Autoencoders. Collapse as an optimization failure, and aggressive inference as the fix.
• Kingma et al. (2016). Improving Variational Inference with Inverse Autoregressive Flow. Introduces IAF and, in appendix C.8, free bits.
• Zhao, Song & Ermon (2017). InfoVAE: Information Maximizing Variational Autoencoders. Reweights the information and mismatch parts of the KL separately, and shows the ELBO, the beta-VAE, and the adversarial autoencoder as members of one family.
• Tolstikhin et al. (2018). Wasserstein Auto-Encoders. The optimal-transport derivation of an aggregate-matching autoencoder, deterministic encoders, and why the IMQ kernel beats the RBF kernel for MMD in code space.
• Locatello et al. (2019). Challenging Common Assumptions in the Unsupervised Learning of Disentangled Representations. The impossibility theorem and the 12,000-model study.
• Child (2021). Very Deep VAEs Generalize Autoregressive Models and Can Outperform Them on Images. Deep hierarchical VAEs without free bits or warm-up, and the appendix figure on architecture and collapse.
• Bowman et al. (2016). Generating Sentences from a Continuous Space. The sentence VAE where collapse was first seen at scale, with KL annealing and word dropout as the fixes.
• Cover & Thomas (2006). Elements of Information Theory, chapter 10. Rate-distortion theory, the Gaussian source, and reverse water-filling, proved carefully.

Watch next

▶

▶

▶